From 759081b1da209e6d41c834f46c881fb714b4bccc Mon Sep 17 00:00:00 2001 From: LHT129 Date: Tue, 7 Jul 2026 14:28:17 +0800 Subject: [PATCH] feat(io): implement CacheIO read-cache layer for disk IO Signed-off-by: LHT129 --- src/inner_string_params.h | 2 + src/io/CMakeLists.txt | 3 + src/io/cache_io/cache_io.h | 189 +++++++++++++++++++++++ src/io/cache_io/cache_io_parameter.cpp | 65 ++++++++ src/io/cache_io/cache_io_parameter.h | 46 ++++++ src/io/cache_io/cache_io_test.cpp | 192 ++++++++++++++++++++++++ src/io/cache_io/lru_page_cache.cpp | 53 +++++++ src/io/cache_io/lru_page_cache.h | 50 ++++++ src/io/cache_io/lru_page_cache_test.cpp | 178 ++++++++++++++++++++++ src/io/cache_io/page.h | 76 ++++++++++ src/io/cache_io/page_cache.cpp | 68 +++++++++ src/io/cache_io/page_cache.h | 93 ++++++++++++ src/io/cache_io/page_cache_test.cpp | 149 ++++++++++++++++++ src/io/common/io_parameter.cpp | 4 + src/io/io_headers.h | 1 + 15 files changed, 1169 insertions(+) create mode 100644 src/io/cache_io/cache_io.h create mode 100644 src/io/cache_io/cache_io_parameter.cpp create mode 100644 src/io/cache_io/cache_io_parameter.h create mode 100644 src/io/cache_io/cache_io_test.cpp create mode 100644 src/io/cache_io/lru_page_cache.cpp create mode 100644 src/io/cache_io/lru_page_cache.h create mode 100644 src/io/cache_io/lru_page_cache_test.cpp create mode 100644 src/io/cache_io/page.h create mode 100644 src/io/cache_io/page_cache.cpp create mode 100644 src/io/cache_io/page_cache.h create mode 100644 src/io/cache_io/page_cache_test.cpp diff --git a/src/inner_string_params.h b/src/inner_string_params.h index e4aff1ea2d..2dd658d229 100644 --- a/src/inner_string_params.h +++ b/src/inner_string_params.h @@ -63,6 +63,7 @@ const char* const IO_TYPE_VALUE_MMAP_IO = "mmap_io"; const char* const IO_TYPE_VALUE_READER_IO = "reader_io"; const char* const IO_TYPE_VALUE_ASYNC_IO = "async_io"; const char* const IO_TYPE_VALUE_BLOCK_MEMORY_IO = "block_memory_io"; +const char* const IO_TYPE_VALUE_CACHE_IO = "cache_io"; const char* const BLOCK_IO_BLOCK_SIZE_KEY = "block_size"; // IO param for file @@ -212,6 +213,7 @@ const std::unordered_map DEFAULT_MAP = { {"IO_TYPE_VALUE_MEMORY_IO", IO_TYPE_VALUE_MEMORY_IO}, {"IO_TYPE_VALUE_BLOCK_MEMORY_IO", IO_TYPE_VALUE_BLOCK_MEMORY_IO}, {"IO_TYPE_VALUE_BUFFER_IO", IO_TYPE_VALUE_BUFFER_IO}, + {"IO_TYPE_VALUE_CACHE_IO", IO_TYPE_VALUE_CACHE_IO}, {"IO_PARAMS_KEY", IO_PARAMS_KEY}, {"BLOCK_IO_BLOCK_SIZE_KEY", BLOCK_IO_BLOCK_SIZE_KEY}, {"QUANTIZATION_TYPE_VALUE_SQ8", QUANTIZATION_TYPE_VALUE_SQ8}, diff --git a/src/io/CMakeLists.txt b/src/io/CMakeLists.txt index fa6dd3203f..3805633c31 100644 --- a/src/io/CMakeLists.txt +++ b/src/io/CMakeLists.txt @@ -28,6 +28,9 @@ set (IO_SRC memory_block_io/memory_block_io.cpp reader_io/reader_io.cpp reader_io/reader_io_parameter.cpp + cache_io/cache_io_parameter.cpp + cache_io/page_cache.cpp + cache_io/lru_page_cache.cpp ) add_library (io OBJECT ${IO_SRC}) diff --git a/src/io/cache_io/cache_io.h b/src/io/cache_io/cache_io.h new file mode 100644 index 0000000000..dc608ecb9e --- /dev/null +++ b/src/io/cache_io/cache_io.h @@ -0,0 +1,189 @@ +// Copyright 2024-present the vsag 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. + +#pragma once + +#include +#include +#include + +#include "io/cache_io/cache_io_parameter.h" +#include "io/cache_io/lru_page_cache.h" +#include "io/cache_io/page_cache.h" +#include "io/common/basic_io.h" + +namespace vsag { + +class IndexCommonParam; + +template +class CacheIOTest; + +/** + * @brief Read-cache IO layer wrapping an inner IO implementation. + * + * CacheIO provides a page-based read cache in front of any disk-backed IO. + * All reads go through the cache; writes passthrough to the inner IO directly. + * The workload is assumed to be non-overlapping between reads and writes. + * + * @tparam IOTmpl The underlying disk IO implementation type. + */ +template +class CacheIO : public BasicIO> { +public: + static constexpr bool InMemory = IOTmpl::InMemory; + static constexpr bool SkipDeserialize = false; + + template + CacheIO(const CacheIOParameter& param, Allocator* allocator, Args&&... args) + : BasicIO>(allocator), + cache_(std::make_unique(std::max( + uint64_t(1), param.total_cache_size_ / std::max(uint64_t(1), param.page_size_)))), + inner_io_(std::make_unique(std::forward(args)...)) { + Page::SetPageSize(std::max(uint64_t(1), param.page_size_)); + this->size_ = inner_io_->size_; + } + + template + CacheIO(uint64_t page_size, uint64_t total_cache_size, Allocator* allocator, Args&&... args) + : BasicIO>(allocator), + cache_(std::make_unique( + std::max(uint64_t(1), total_cache_size / std::max(uint64_t(1), page_size)))), + inner_io_(std::make_unique(std::forward(args)...)) { + Page::SetPageSize(std::max(uint64_t(1), page_size)); + this->size_ = inner_io_->size_; + } + + ~CacheIO() = default; + + void + WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset) { + inner_io_->WriteImpl(data, size, offset); + if (offset + size > this->size_) { + this->size_ = offset + size; + } + } + + bool + ReadImpl(uint64_t size, uint64_t offset, uint8_t* data) const { + if (not this->check_valid_offset(size + offset)) { + return false; + } + uint64_t page_size = Page::PageSize(); + uint64_t cur_size = 0; + while (cur_size < size) { + uint64_t cur_offset = offset + cur_size; + uint64_t page_id = cur_offset / page_size; + uint64_t page_offset = cur_offset % page_size; + uint64_t copy_size = std::min(size - cur_size, page_size - page_offset); + + PagePtr page = get_or_load_page(page_id); + if (page == nullptr) { + return false; + } + std::memcpy(data + cur_size, page->Data() + page_offset, copy_size); + cur_size += copy_size; + } + return true; + } + + [[nodiscard]] const uint8_t* + DirectReadImpl(uint64_t size, uint64_t offset, bool& need_release) const { + need_release = false; + if (size == 0 or not this->check_valid_offset(size + offset)) { + return nullptr; + } + auto* data = reinterpret_cast(this->allocator_->Allocate(size)); + if (data == nullptr) { + return nullptr; + } + if (not ReadImpl(size, offset, data)) { + this->allocator_->Deallocate(data); + return nullptr; + } + need_release = true; + return data; + } + + bool + MultiReadImpl(uint8_t* datas, uint64_t* sizes, uint64_t* offsets, uint64_t count) const { + for (uint64_t i = 0; i < count; ++i) { + if (not this->ReadImpl(sizes[i], offsets[i], datas)) { + return false; + } + datas += sizes[i]; + } + return true; + } + + void + ReleaseImpl(const uint8_t* data) const { + this->allocator_->Deallocate(const_cast(data)); + } + + void + InitIOImpl(const IOParamPtr& io_param) { + inner_io_->InitIO(io_param); + this->size_ = inner_io_->size_; + } + + void + ResizeImpl(uint64_t size) { + inner_io_->ResizeImpl(size); + this->size_ = size; + } + +private: + /** + * @brief Get a page from the cache, loading it from the inner IO on miss. + * + * The returned PagePtr keeps the page alive even if it is evicted + * from the cache concurrently. + * + * @param page_id The page to fetch. + * @return The page, or nullptr on failure. + */ + PagePtr + get_or_load_page(uint64_t page_id) const { + PagePtr page = cache_->Get(page_id); + if (page != nullptr) { + return page; + } + + uint64_t page_size = Page::PageSize(); + uint64_t file_offset = page_id * page_size; + if (file_offset >= this->size_) { + return nullptr; + } + uint64_t read_size = std::min(page_size, this->size_ - file_offset); + + auto new_page = std::make_shared(this->allocator_); + if (new_page->Data() == nullptr) { + return nullptr; + } + + if (not inner_io_->ReadImpl(read_size, file_offset, new_page->Data())) { + return nullptr; + } + return cache_->Insert(page_id, std::move(new_page)); + } + +private: + friend class CacheIOTest; + + std::unique_ptr cache_; + std::unique_ptr inner_io_; +}; + +} // namespace vsag diff --git a/src/io/cache_io/cache_io_parameter.cpp b/src/io/cache_io/cache_io_parameter.cpp new file mode 100644 index 0000000000..3cc26b3678 --- /dev/null +++ b/src/io/cache_io/cache_io_parameter.cpp @@ -0,0 +1,65 @@ +// Copyright 2024-present the vsag 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 "io/cache_io/cache_io_parameter.h" + +#include + +#include "inner_string_params.h" +#include "io/cache_io/page.h" + +namespace vsag { + +static const char* const CACHE_IO_PAGE_SIZE_KEY = "page_size"; +static const char* const CACHE_IO_TOTAL_CACHE_SIZE_KEY = "total_cache_size"; +static const char* const CACHE_IO_EVICTION_STRATEGY_KEY = "eviction_strategy"; + +static constexpr uint64_t DEFAULT_PAGE_SIZE = Page::DEFAULT_PAGE_SIZE; + +CacheIOParameter::CacheIOParameter() : IOParameter(IO_TYPE_VALUE_CACHE_IO) { +} + +CacheIOParameter::CacheIOParameter(const vsag::JsonType& json) + : IOParameter(IO_TYPE_VALUE_CACHE_IO) { + this->FromJson(json); // NOLINT(clang-analyzer-optin.cplusplus.VirtualCall) +} + +void +CacheIOParameter::FromJson(const JsonType& json) { + if (json.Contains(CACHE_IO_PAGE_SIZE_KEY)) { + uint64_t parsed = json[CACHE_IO_PAGE_SIZE_KEY].GetUint64(); + this->page_size_ = parsed > 0 ? parsed : DEFAULT_PAGE_SIZE; + } + if (json.Contains(CACHE_IO_TOTAL_CACHE_SIZE_KEY)) { + this->total_cache_size_ = json[CACHE_IO_TOTAL_CACHE_SIZE_KEY].GetUint64(); + } + if (json.Contains(CACHE_IO_EVICTION_STRATEGY_KEY)) { + this->eviction_strategy_ = json[CACHE_IO_EVICTION_STRATEGY_KEY].GetString(); + if (this->eviction_strategy_ != "lru") { + this->eviction_strategy_ = "lru"; + } + } +} + +JsonType +CacheIOParameter::ToJson() const { + JsonType json; + json[TYPE_KEY].SetString(IO_TYPE_VALUE_CACHE_IO); + json[CACHE_IO_PAGE_SIZE_KEY].SetUint64(this->page_size_); + json[CACHE_IO_TOTAL_CACHE_SIZE_KEY].SetUint64(this->total_cache_size_); + json[CACHE_IO_EVICTION_STRATEGY_KEY].SetString(this->eviction_strategy_); + return json; +} + +} // namespace vsag diff --git a/src/io/cache_io/cache_io_parameter.h b/src/io/cache_io/cache_io_parameter.h new file mode 100644 index 0000000000..fbf9872479 --- /dev/null +++ b/src/io/cache_io/cache_io_parameter.h @@ -0,0 +1,46 @@ + +// Copyright 2024-present the vsag 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. + +#pragma once + +#include +#include + +#include "io/common/io_parameter.h" +#include "utils/pointer_define.h" + +namespace vsag { + +DEFINE_POINTER(CacheIOParameter); + +class CacheIOParameter : public IOParameter { +public: + CacheIOParameter(); + + explicit CacheIOParameter(const JsonType& json); + + void + FromJson(const JsonType& json) override; + + JsonType + ToJson() const override; + +public: + uint64_t page_size_{131072}; + uint64_t total_cache_size_{268435456}; + std::string eviction_strategy_{"lru"}; +}; + +} // namespace vsag diff --git a/src/io/cache_io/cache_io_test.cpp b/src/io/cache_io/cache_io_test.cpp new file mode 100644 index 0000000000..38d05be0f7 --- /dev/null +++ b/src/io/cache_io/cache_io_test.cpp @@ -0,0 +1,192 @@ +// Copyright 2024-present the vsag 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 "io/cache_io/cache_io.h" + +#include +#include + +#include "impl/allocator/safe_allocator.h" +#include "io/buffer_io/buffer_io.h" +#include "io/cache_io/cache_io_parameter.h" +#include "io/common/basic_io_test.h" +#include "io/memory_io/memory_io.h" +#include "io/mmap_io/mmap_io.h" +#include "unittest.h" + +using namespace vsag; + +template +void +TestCacheIOFileBackend(const std::string& path) { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CacheIOParameter param; + param.page_size_ = 128; + param.total_cache_size_ = 128 * 4; + + CacheIO io(param, allocator.get(), path, allocator.get()); + TestBasicReadWrite(io); + + std::vector data(300); + for (uint64_t i = 0; i < data.size(); ++i) { + data[i] = static_cast(i % 251); + } + io.Write(data.data(), data.size(), 32); + + std::vector read_buf(data.size()); + REQUIRE(io.Read(data.size(), 32, read_buf.data())); + REQUIRE(read_buf == data); +} + +TEST_CASE("CacheIO Basic Test", "[CacheIO][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CacheIOParameter param; + param.page_size_ = 128; + param.total_cache_size_ = 128 * 8; + + CacheIO io(param, allocator.get(), allocator.get()); + TestBasicReadWrite(io); +} + +TEST_CASE("CacheIO BufferIO Backend Test", "[CacheIO][ut]") { + fixtures::TempDir dir("cache_io_buffer"); + TestCacheIOFileBackend(dir.GenerateRandomFile(false)); +} + +TEST_CASE("CacheIO MMapIO Backend Test", "[CacheIO][ut]") { + fixtures::TempDir dir("cache_io_mmap"); + TestCacheIOFileBackend(dir.GenerateRandomFile(false)); +} + +TEST_CASE("CacheIO Eviction Test", "[CacheIO][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CacheIOParameter param; + param.page_size_ = 64; + param.total_cache_size_ = 64 * 2; + + CacheIO io(param, allocator.get(), allocator.get()); + + std::vector data(64, 0xAA); + io.Write(data.data(), 64, 0); + std::vector data2(64, 0xBB); + io.Write(data2.data(), 64, 64); + std::vector data3(64, 0xCC); + io.Write(data3.data(), 64, 128); + std::vector data4(64, 0xDD); + io.Write(data4.data(), 64, 192); + + std::vector read_buf(64); + io.Read(64, 192, read_buf.data()); + REQUIRE(memcmp(read_buf.data(), data4.data(), 64) == 0); + + io.Read(64, 128, read_buf.data()); + REQUIRE(memcmp(read_buf.data(), data3.data(), 64) == 0); + + io.Read(64, 0, read_buf.data()); + REQUIRE(memcmp(read_buf.data(), data.data(), 64) == 0); + + io.Read(64, 64, read_buf.data()); + REQUIRE(memcmp(read_buf.data(), data2.data(), 64) == 0); +} + +TEST_CASE("CacheIO DirectRead Single Page Test", "[CacheIO][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CacheIOParameter param; + param.page_size_ = 256; + param.total_cache_size_ = 256 * 4; + + CacheIO io(param, allocator.get(), allocator.get()); + + std::vector data(100, 0xAB); + io.Write(data.data(), 100, 50); + + bool need_release = false; + const auto* ptr = io.Read(100, 50, need_release); + REQUIRE(ptr != nullptr); + REQUIRE(need_release == true); + REQUIRE(memcmp(ptr, data.data(), 100) == 0); + io.Release(ptr); +} + +TEST_CASE("CacheIO DirectRead Cross Page Test", "[CacheIO][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CacheIOParameter param; + param.page_size_ = 64; + param.total_cache_size_ = 64 * 4; + + CacheIO io(param, allocator.get(), allocator.get()); + + std::vector data(100, 0xCD); + io.Write(data.data(), 100, 30); + + bool need_release = false; + const auto* ptr = io.Read(100, 30, need_release); + REQUIRE(ptr != nullptr); + REQUIRE(need_release == true); + REQUIRE(memcmp(ptr, data.data(), 100) == 0); + io.Release(ptr); +} + +TEST_CASE("CacheIO MultiRead Test", "[CacheIO][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CacheIOParameter param; + param.page_size_ = 64; + param.total_cache_size_ = 64 * 4; + + CacheIO io(param, allocator.get(), allocator.get()); + + std::vector data1(30, 0x11); + std::vector data2(40, 0x22); + io.Write(data1.data(), 30, 10); + io.Write(data2.data(), 40, 100); + + uint64_t sizes[] = {30, 40}; + uint64_t offsets[] = {10, 100}; + std::vector result(70); + bool ret = io.MultiRead(result.data(), sizes, offsets, 2); + REQUIRE(ret); + REQUIRE(memcmp(result.data(), data1.data(), 30) == 0); + REQUIRE(memcmp(result.data() + 30, data2.data(), 40) == 0); +} + +TEST_CASE("CacheIO Parameter Test", "[CacheIO][ut]") { + CacheIOParameter param; + REQUIRE(param.page_size_ == 131072); + REQUIRE(param.total_cache_size_ == 268435456); + REQUIRE(param.eviction_strategy_ == "lru"); + + JsonType json; + json["type"].SetString("cache_io"); + json["page_size"].SetUint64(1024); + json["total_cache_size"].SetUint64(4096); + json["eviction_strategy"].SetString("lru"); + + CacheIOParameter param2(json); + REQUIRE(param2.page_size_ == 1024); + REQUIRE(param2.total_cache_size_ == 4096); + REQUIRE(param2.eviction_strategy_ == "lru"); + + auto json2 = param2.ToJson(); + REQUIRE(json2["page_size"].GetUint64() == 1024); + REQUIRE(json2["total_cache_size"].GetUint64() == 4096); +} + +TEST_CASE("CacheIO Parameter Validation Test", "[CacheIO][ut]") { + JsonType json; + json["type"].SetString("cache_io"); + json["page_size"].SetUint64(0); + + CacheIOParameter param(json); + REQUIRE(param.page_size_ == 131072); +} diff --git a/src/io/cache_io/lru_page_cache.cpp b/src/io/cache_io/lru_page_cache.cpp new file mode 100644 index 0000000000..5f3710db66 --- /dev/null +++ b/src/io/cache_io/lru_page_cache.cpp @@ -0,0 +1,53 @@ +// Copyright 2024-present the vsag 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 "io/cache_io/lru_page_cache.h" + +namespace vsag { + +LRUPageCache::LRUPageCache(uint64_t max_pages) : PageCache(max_pages) { +} + +void +LRUPageCache::OnAccess(uint64_t page_id) { + auto it = iters_.find(page_id); + if (it != iters_.end()) { + order_.splice(order_.begin(), order_, it->second); + } +} + +void +LRUPageCache::OnInsert(uint64_t page_id) { + order_.push_front(page_id); + iters_[page_id] = order_.begin(); +} + +void +LRUPageCache::OnRemove(uint64_t page_id) { + auto it = iters_.find(page_id); + if (it != iters_.end()) { + order_.erase(it->second); + iters_.erase(it); + } +} + +uint64_t +LRUPageCache::PickVictim() { + if (order_.empty()) { + return UINT64_MAX; + } + return order_.back(); +} + +} // namespace vsag diff --git a/src/io/cache_io/lru_page_cache.h b/src/io/cache_io/lru_page_cache.h new file mode 100644 index 0000000000..15e918d5d1 --- /dev/null +++ b/src/io/cache_io/lru_page_cache.h @@ -0,0 +1,50 @@ +// Copyright 2024-present the vsag 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. + +#pragma once + +#include +#include +#include + +#include "io/cache_io/page_cache.h" + +namespace vsag { + +/** + * @brief PageCache with Least-Recently-Used eviction policy. + */ +class LRUPageCache : public PageCache { +public: + explicit LRUPageCache(uint64_t max_pages); + +protected: + void + OnAccess(uint64_t page_id) override; + + void + OnInsert(uint64_t page_id) override; + + void + OnRemove(uint64_t page_id) override; + + uint64_t + PickVictim() override; + +private: + std::list order_; + std::unordered_map::iterator> iters_; +}; + +} // namespace vsag diff --git a/src/io/cache_io/lru_page_cache_test.cpp b/src/io/cache_io/lru_page_cache_test.cpp new file mode 100644 index 0000000000..6eaa2fc041 --- /dev/null +++ b/src/io/cache_io/lru_page_cache_test.cpp @@ -0,0 +1,178 @@ +// Copyright 2024-present the vsag 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 "io/cache_io/lru_page_cache.h" + +#include + +#include "impl/allocator/safe_allocator.h" +#include "unittest.h" + +using namespace vsag; + +namespace { + +PagePtr +MakePage(Allocator* allocator) { + auto page = std::make_shared(allocator); + if (page->Data() != nullptr) { + std::memset(page->Data(), 0xAA, Page::PageSize()); + } + return page; +} + +} // namespace + +TEST_CASE("LRUPageCache Basic Insert and Get", "[LRUPageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + LRUPageCache cache(4); + + REQUIRE(cache.Get(0) == nullptr); + REQUIRE(cache.Size() == 0); + + auto page0 = MakePage(allocator.get()); + cache.Insert(0, page0); + REQUIRE(cache.Size() == 1); + REQUIRE(cache.Get(0) == page0); +} + +TEST_CASE("LRUPageCache Insert Duplicate Returns Existing", "[LRUPageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + LRUPageCache cache(4); + + auto page0 = MakePage(allocator.get()); + cache.Insert(0, page0); + + auto page0_new = MakePage(allocator.get()); + auto result = cache.Insert(0, page0_new); + REQUIRE(result == page0); + REQUIRE(cache.Size() == 1); +} + +TEST_CASE("LRUPageCache LRU Eviction Order", "[LRUPageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + LRUPageCache cache(3); + + auto page0 = MakePage(allocator.get()); + auto page1 = MakePage(allocator.get()); + auto page2 = MakePage(allocator.get()); + auto page3 = MakePage(allocator.get()); + + cache.Insert(0, page0); + cache.Insert(1, page1); + cache.Insert(2, page2); + // order: 2 -> 1 -> 0 + + // Access page0, moving it to front: 0 -> 2 -> 1 + cache.Get(0); + + // Insert page3, should evict page1 (LRU): 3 -> 0 -> 2 + cache.Insert(3, page3); + REQUIRE(cache.Size() == 3); + REQUIRE(cache.Get(1) == nullptr); + REQUIRE(cache.Get(0) != nullptr); + REQUIRE(cache.Get(2) != nullptr); + REQUIRE(cache.Get(3) != nullptr); +} + +TEST_CASE("LRUPageCache Capacity One", "[LRUPageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + LRUPageCache cache(1); + + auto page0 = MakePage(allocator.get()); + auto page1 = MakePage(allocator.get()); + + cache.Insert(0, page0); + REQUIRE(cache.Size() == 1); + REQUIRE(cache.Get(0) == page0); + + cache.Insert(1, page1); + REQUIRE(cache.Size() == 1); + REQUIRE(cache.Get(0) == nullptr); + REQUIRE(cache.Get(1) == page1); +} + +TEST_CASE("LRUPageCache Remove Updates LRU Order", "[LRUPageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + LRUPageCache cache(3); + + auto page0 = MakePage(allocator.get()); + auto page1 = MakePage(allocator.get()); + + cache.Insert(0, page0); + cache.Insert(1, page1); + cache.Remove(0); + REQUIRE(cache.Size() == 1); + REQUIRE(cache.Get(0) == nullptr); + + auto page2 = MakePage(allocator.get()); + cache.Insert(2, page2); + REQUIRE(cache.Size() == 2); + REQUIRE(cache.Get(1) == page1); + REQUIRE(cache.Get(2) == page2); +} + +TEST_CASE("LRUPageCache Remove Nonexistent Is No-op", "[LRUPageCache][ut]") { + LRUPageCache cache(4); + REQUIRE_NOTHROW(cache.Remove(999)); + REQUIRE(cache.Size() == 0); +} + +TEST_CASE("LRUPageCache Size After Multiple Evictions", "[LRUPageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + LRUPageCache cache(2); + + cache.Insert(0, MakePage(allocator.get())); + cache.Insert(1, MakePage(allocator.get())); + cache.Insert(2, MakePage(allocator.get())); + REQUIRE(cache.Size() == 2); + cache.Insert(3, MakePage(allocator.get())); + REQUIRE(cache.Size() == 2); + cache.Insert(4, MakePage(allocator.get())); + REQUIRE(cache.Size() == 2); +} + +TEST_CASE("LRUPageCache Get Updates Recency Without Removing", "[LRUPageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + LRUPageCache cache(3); + + auto page0 = MakePage(allocator.get()); + auto page1 = MakePage(allocator.get()); + auto page2 = MakePage(allocator.get()); + + cache.Insert(0, page0); + cache.Insert(1, page1); + cache.Insert(2, page2); + // order: 2 -> 1 -> 0 + + cache.Get(1); + // order: 1 -> 2 -> 0 + + cache.Get(0); + // order: 0 -> 1 -> 2 + + cache.Insert(3, MakePage(allocator.get())); + // evicts page2 (LRU): 3 -> 0 -> 1 + REQUIRE(cache.Get(2) == nullptr); + REQUIRE(cache.Get(0) != nullptr); + REQUIRE(cache.Get(1) != nullptr); + REQUIRE(cache.Get(3) != nullptr); +} diff --git a/src/io/cache_io/page.h b/src/io/cache_io/page.h new file mode 100644 index 0000000000..a8fad642bc --- /dev/null +++ b/src/io/cache_io/page.h @@ -0,0 +1,76 @@ +// Copyright 2024-present the vsag 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. + +#pragma once + +#include +#include +#include +#include + +#include "vsag/allocator.h" + +namespace vsag { + +/** + * @brief A cached page owning its data buffer via RAII. + * + * All pages share a process-wide uniform page size. + */ +class Page { +public: + explicit Page(Allocator* allocator) : allocator_(allocator) { + data_ = static_cast(allocator_->Allocate(page_size_)); + } + + ~Page() { + if (data_ != nullptr) { + allocator_->Deallocate(data_); + } + } + + Page(const Page&) = delete; + Page& + operator=(const Page&) = delete; + + [[nodiscard]] uint8_t* + Data() const { + return data_; + } + + static void + SetPageSize(uint64_t page_size) { + static std::mutex mtx; + std::lock_guard lock(mtx); + page_size_ = page_size; + } + + static uint64_t + PageSize() { + return page_size_; + } + +public: + static constexpr uint64_t DEFAULT_PAGE_SIZE = 128 * 1024; + +private: + inline static std::atomic page_size_{DEFAULT_PAGE_SIZE}; + + uint8_t* data_{nullptr}; + Allocator* allocator_{nullptr}; +}; + +using PagePtr = std::shared_ptr; + +} // namespace vsag diff --git a/src/io/cache_io/page_cache.cpp b/src/io/cache_io/page_cache.cpp new file mode 100644 index 0000000000..85e9241019 --- /dev/null +++ b/src/io/cache_io/page_cache.cpp @@ -0,0 +1,68 @@ +// Copyright 2024-present the vsag 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 "io/cache_io/page_cache.h" + +namespace vsag { + +PageCache::PageCache(uint64_t max_pages) : max_pages_(max_pages) { +} + +PagePtr +PageCache::Get(uint64_t page_id) { + std::lock_guard lock(mutex_); + auto it = pages_.find(page_id); + if (it == pages_.end()) { + return nullptr; + } + OnAccess(page_id); + return it->second; +} + +PagePtr +PageCache::Insert(uint64_t page_id, PagePtr page) { + std::lock_guard lock(mutex_); + auto existing = pages_.find(page_id); + if (existing != pages_.end()) { + OnAccess(page_id); + return existing->second; + } + while (pages_.size() >= max_pages_) { + uint64_t victim = PickVictim(); + if (victim == UINT64_MAX) { + break; + } + pages_.erase(victim); + OnRemove(victim); + } + pages_[page_id] = std::move(page); + OnInsert(page_id); + return pages_[page_id]; +} + +void +PageCache::Remove(uint64_t page_id) { + std::lock_guard lock(mutex_); + if (pages_.erase(page_id) > 0) { + OnRemove(page_id); + } +} + +uint64_t +PageCache::Size() const { + std::lock_guard lock(mutex_); + return pages_.size(); +} + +} // namespace vsag diff --git a/src/io/cache_io/page_cache.h b/src/io/cache_io/page_cache.h new file mode 100644 index 0000000000..c1bdd9efbc --- /dev/null +++ b/src/io/cache_io/page_cache.h @@ -0,0 +1,93 @@ +// Copyright 2024-present the vsag 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. + +#pragma once + +#include +#include +#include + +#include "io/cache_io/page.h" + +namespace vsag { + +/** + * @brief Page storage manager with pluggable eviction policy. + * + * PageCache owns the cached pages and decides which page to evict + * when the capacity is reached. Pages are managed by shared_ptr: + * a page returned to the caller stays valid even if it is evicted + * from the cache concurrently. + */ +class PageCache { +public: + explicit PageCache(uint64_t max_pages); + + virtual ~PageCache() = default; + + /** + * @brief Look up a page, marking it as recently accessed. + * + * @param page_id The page to look up. + * @return The cached page, or nullptr on miss. + */ + PagePtr + Get(uint64_t page_id); + + /** + * @brief Insert a page, evicting victims first if the cache is full. + * + * If the page already exists, the existing page is returned instead. + * + * @param page_id The page id. + * @param page The page to insert. + * @return The page stored in the cache. + */ + PagePtr + Insert(uint64_t page_id, PagePtr page); + + /** + * @brief Remove a page from the cache. + * + * @param page_id The page to remove. + */ + void + Remove(uint64_t page_id); + + /** + * @brief Number of pages currently cached. + */ + uint64_t + Size() const; + +protected: + virtual void + OnAccess(uint64_t page_id) = 0; + + virtual void + OnInsert(uint64_t page_id) = 0; + + virtual void + OnRemove(uint64_t page_id) = 0; + + virtual uint64_t + PickVictim() = 0; + +protected: + mutable std::mutex mutex_; + std::unordered_map pages_; + uint64_t max_pages_{0}; +}; + +} // namespace vsag diff --git a/src/io/cache_io/page_cache_test.cpp b/src/io/cache_io/page_cache_test.cpp new file mode 100644 index 0000000000..e609020816 --- /dev/null +++ b/src/io/cache_io/page_cache_test.cpp @@ -0,0 +1,149 @@ +// Copyright 2024-present the vsag 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 "io/cache_io/page_cache.h" + +#include +#include + +#include "impl/allocator/safe_allocator.h" +#include "unittest.h" + +using namespace vsag; + +namespace { + +class TestPageCache : public PageCache { +public: + explicit TestPageCache(uint64_t max_pages) : PageCache(max_pages) { + } + +protected: + void + OnAccess(uint64_t /*page_id*/) override { + } + + void + OnInsert(uint64_t /*page_id*/) override { + } + + void + OnRemove(uint64_t /*page_id*/) override { + } + + uint64_t + PickVictim() override { + if (pages_.empty()) { + return UINT64_MAX; + } + return pages_.begin()->first; + } +}; + +PagePtr +MakePage(Allocator* allocator) { + auto page = std::make_shared(allocator); + if (page->Data() != nullptr) { + std::memset(page->Data(), 0xAA, Page::PageSize()); + } + return page; +} + +} // namespace + +TEST_CASE("PageCache Insert and Get", "[PageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + TestPageCache cache(4); + + REQUIRE(cache.Size() == 0); + REQUIRE(cache.Get(0) == nullptr); + + auto page0 = MakePage(allocator.get()); + auto result = cache.Insert(0, page0); + REQUIRE(result != nullptr); + REQUIRE(cache.Size() == 1); + + auto fetched = cache.Get(0); + REQUIRE(fetched == page0); +} + +TEST_CASE("PageCache Insert Duplicate Returns Existing", "[PageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + TestPageCache cache(4); + + auto page0 = MakePage(allocator.get()); + cache.Insert(0, page0); + + auto page0_new = MakePage(allocator.get()); + auto result = cache.Insert(0, page0_new); + REQUIRE(result == page0); + REQUIRE(result != page0_new); + REQUIRE(cache.Size() == 1); +} + +TEST_CASE("PageCache Eviction When Full", "[PageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + TestPageCache cache(2); + + auto page0 = MakePage(allocator.get()); + auto page1 = MakePage(allocator.get()); + auto page2 = MakePage(allocator.get()); + + cache.Insert(0, page0); + cache.Insert(1, page1); + REQUIRE(cache.Size() == 2); + + cache.Insert(2, page2); + REQUIRE(cache.Size() == 2); + // TestPageCache evicts first map entry (unspecified order), + // so we just verify one page was evicted. +} + +TEST_CASE("PageCache Remove", "[PageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + TestPageCache cache(4); + + auto page0 = MakePage(allocator.get()); + cache.Insert(0, page0); + REQUIRE(cache.Size() == 1); + + cache.Remove(0); + REQUIRE(cache.Size() == 0); + REQUIRE(cache.Get(0) == nullptr); +} + +TEST_CASE("PageCache Remove Nonexistent", "[PageCache][ut]") { + TestPageCache cache(4); + REQUIRE_NOTHROW(cache.Remove(999)); + REQUIRE(cache.Size() == 0); +} + +TEST_CASE("PageCache Size Tracking", "[PageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + Page::SetPageSize(64); + TestPageCache cache(8); + + for (uint64_t i = 0; i < 5; ++i) { + cache.Insert(i, MakePage(allocator.get())); + } + REQUIRE(cache.Size() == 5); + + cache.Remove(1); + cache.Remove(3); + REQUIRE(cache.Size() == 3); +} diff --git a/src/io/common/io_parameter.cpp b/src/io/common/io_parameter.cpp index b38060d517..fc85f28095 100644 --- a/src/io/common/io_parameter.cpp +++ b/src/io/common/io_parameter.cpp @@ -21,6 +21,7 @@ #include "inner_string_params.h" #include "io/async_io/async_io_parameter.h" #include "io/buffer_io/buffer_io_parameter.h" +#include "io/cache_io/cache_io_parameter.h" #include "io/memory_block_io/memory_block_io_parameter.h" #include "io/memory_io/memory_io_parameter.h" #include "io/mmap_io/mmap_io_parameter.h" @@ -59,6 +60,9 @@ IOParameter::GetIOParameterByJson(const JsonType& json) { } else if (type_name == IO_TYPE_VALUE_MMAP_IO) { io_ptr = std::make_shared(); io_ptr->FromJson(json); + } else if (type_name == IO_TYPE_VALUE_CACHE_IO) { + io_ptr = std::make_shared(); + io_ptr->FromJson(json); } else if (type_name == IO_TYPE_VALUE_READER_IO) { io_ptr = std::make_shared(); io_ptr->FromJson(json); diff --git a/src/io/io_headers.h b/src/io/io_headers.h index bf05377a37..40fb826dd3 100644 --- a/src/io/io_headers.h +++ b/src/io/io_headers.h @@ -17,6 +17,7 @@ #include "io/async_io/async_io.h" #include "io/buffer_io/buffer_io.h" +#include "io/cache_io/cache_io.h" #include "io/common/basic_io.h" #include "io/memory_block_io/memory_block_io.h" #include "io/memory_io/memory_io.h"