feat(io): implement CacheIO read-cache layer for disk IO

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<std::string, std::string> 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},

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})

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 <algorithm>
+#include <cstring>
+#include <memory>
+
+#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 <typename IOTmpl>
+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 <typename IOTmpl>
+class CacheIO : public BasicIO<CacheIO<IOTmpl>> {
+public:
+    static constexpr bool InMemory = IOTmpl::InMemory;
+    static constexpr bool SkipDeserialize = false;
+
+    template <typename... Args>
+    CacheIO(const CacheIOParameter& param, Allocator* allocator, Args&&... args)
+        : BasicIO<CacheIO<IOTmpl>>(allocator),
+          cache_(std::make_unique<LRUPageCache>(std::max(
+              uint64_t(1), param.total_cache_size_ / std::max(uint64_t(1), param.page_size_)))),
+          inner_io_(std::make_unique<IOTmpl>(std::forward<Args>(args)...)) {
+        Page::SetPageSize(std::max(uint64_t(1), param.page_size_));
+        this->size_ = inner_io_->size_;
+    }
+
+    template <typename... Args>
+    CacheIO(uint64_t page_size, uint64_t total_cache_size, Allocator* allocator, Args&&... args)
+        : BasicIO<CacheIO<IOTmpl>>(allocator),
+          cache_(std::make_unique<LRUPageCache>(
+              std::max(uint64_t(1), total_cache_size / std::max(uint64_t(1), page_size)))),
+          inner_io_(std::make_unique<IOTmpl>(std::forward<Args>(args)...)) {
+        Page::SetPageSize(std::max(uint64_t(1), page_size));
+        this->size_ = inner_io_->size_;
+    }
+
+    ~CacheIO() override = 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<uint8_t*>(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<uint8_t*>(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<Page>(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<IOTmpl>;
+
+    std::unique_ptr<PageCache> cache_;
+    std::unique_ptr<IOTmpl> inner_io_;
+};
+
+}  // namespace vsag

src/io/cache_io/cache_io_parameter.cpp

@@ -0,0 +1,60 @@
+// 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 "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();
+    }
+}
+
+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

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 <cstdint>
+#include <string>
+
+#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