introduce simple resource_pool

- for visited_list, give an implement of visited_list_pool - we will have some other object like aio_context... - currently hgraph use visitlist from hnswlib, now make it global Signed-off-by: LHT129 <[email protected]>
alipay · Nov 12, 2024 · d2993eb · d2993eb
1 parent 50d2538
commit d2993eb
Show file tree

Hide file tree

Showing 5 changed files with 368 additions and 23 deletions.
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -1,27 +1,30 @@
 
-add_subdirectory (simd)
+add_subdirectory(simd)
 
-file (GLOB CPP_SRCS "*.cpp")
-file (GLOB CPP_FACTORY_SRCS "factory/*.cpp")
-file (GLOB CPP_CONJUGATE_GRAPH_SRCS "impl/*.cpp")
-file (GLOB CPP_INDEX_SRCS "index/*.cpp")
-file (GLOB CPP_HNSWLIB_SRCS "algorithm/hnswlib/*.cpp")
-file (GLOB CPP_QUANTIZATION_SRCS "quantization/*.cpp")
-file (GLOB CPP_DATA_CELL_SRCS "data_cell/*.cpp")
-list (FILTER CPP_SRCS EXCLUDE REGEX "_test.cpp")
-list (FILTER CPP_FACTORY_SRCS EXCLUDE REGEX "_test.cpp")
-list (FILTER CPP_CONJUGATE_GRAPH_SRCS EXCLUDE REGEX "_test.cpp")
-list (FILTER CPP_INDEX_SRCS EXCLUDE REGEX "_test.cpp")
-list (FILTER CPP_QUANTIZATION_SRCS EXCLUDE REGEX "_test.cpp")
-list (FILTER CPP_DATA_CELL_SRCS EXCLUDE REGEX "_test.cpp")
+file(GLOB CPP_SRCS "*.cpp")
+file(GLOB CPP_FACTORY_SRCS "factory/*.cpp")
+file(GLOB CPP_CONJUGATE_GRAPH_SRCS "impl/*.cpp")
+file(GLOB CPP_INDEX_SRCS "index/*.cpp")
+file(GLOB CPP_HNSWLIB_SRCS "algorithm/hnswlib/*.cpp")
+file(GLOB CPP_QUANTIZATION_SRCS "quantization/*.cpp")
+file(GLOB CPP_DATA_CELL_SRCS "data_cell/*.cpp")
+list(FILTER CPP_SRCS EXCLUDE REGEX "_test.cpp")
+list(FILTER CPP_FACTORY_SRCS EXCLUDE REGEX "_test.cpp")
+list(FILTER CPP_CONJUGATE_GRAPH_SRCS EXCLUDE REGEX "_test.cpp")
+list(FILTER CPP_INDEX_SRCS EXCLUDE REGEX "_test.cpp")
+list(FILTER CPP_QUANTIZATION_SRCS EXCLUDE REGEX "_test.cpp")
+list(FILTER CPP_DATA_CELL_SRCS EXCLUDE REGEX "_test.cpp")
 
-set (VSAG_SRCS ${CPP_SRCS} ${CPP_FACTORY_SRCS} ${CPP_INDEX_SRCS}
-        ${CPP_CONJUGATE_GRAPH_SRCS} ${CPP_HNSWLIB_SRCS} ${CPP_QUANTIZATION_SRCS} ${CPP_DATA_CELL_SRCS})
-add_library (vsag SHARED ${VSAG_SRCS})
-add_library (vsag_static STATIC ${VSAG_SRCS})
+file(GLOB CPP_UTILS_SRCS "*.cpp")
+list(FILTER CPP_UTILS_SRCS EXCLUDE REGEX "_test.cpp")
 
-set (VSAG_DEP_LIBS diskann pthread m dl simd fmt::fmt-header-only nlohmann_json::nlohmann_json roaring)
-target_link_libraries (vsag ${VSAG_DEP_LIBS})
-target_link_libraries (vsag_static ${VSAG_DEP_LIBS})
-maybe_add_dependencies (vsag spdlog roaring openblas boost mkl)
-maybe_add_dependencies (vsag_static spdlog roaring openblas boost mkl)
+set(VSAG_SRCS ${CPP_SRCS} ${CPP_FACTORY_SRCS} ${CPP_INDEX_SRCS}
+        ${CPP_CONJUGATE_GRAPH_SRCS} ${CPP_HNSWLIB_SRCS} ${CPP_QUANTIZATION_SRCS} ${CPP_DATA_CELL_SRCS} ${CPP_UTILS_SRCS})
+add_library(vsag SHARED ${VSAG_SRCS})
+add_library(vsag_static STATIC ${VSAG_SRCS})
+
+set(VSAG_DEP_LIBS diskann pthread m dl simd fmt::fmt-header-only nlohmann_json::nlohmann_json roaring)
+target_link_libraries(vsag ${VSAG_DEP_LIBS})
+target_link_libraries(vsag_static ${VSAG_DEP_LIBS})
+maybe_add_dependencies(vsag spdlog roaring openblas boost mkl)
+maybe_add_dependencies(vsag_static spdlog roaring openblas boost mkl)
diff --git a/src/utils/resource_object.h b/src/utils/resource_object.h
@@ -0,0 +1,30 @@
+
+// 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
+
+namespace vsag {
+
+class ResourceObject {
+public:
+    ResourceObject() = default;
+
+    virtual ~ResourceObject() = default;
+
+    virtual void
+    Reset() = 0;
+};
+
+}  // namespace vsag
diff --git a/src/utils/resource_object_pool.h b/src/utils/resource_object_pool.h
@@ -0,0 +1,99 @@
+
+// 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 <functional>
+#include <memory>
+#include <mutex>
+#include <type_traits>
+
+#include "resource_object.h"
+namespace vsag {
+
+template <typename T,
+          typename = typename std::enable_if<std::is_base_of<ResourceObject, T>::value>::type>
+class ResourceObjectPool {
+public:
+    using ConstructFuncType = std::function<std::shared_ptr<T>()>;
+
+public:
+    template <typename... Args>
+    explicit ResourceObjectPool(uint64_t init_size, Args... args) {
+        this->pool_size_ = init_size;
+        this->constructor_ = [=]() -> std::shared_ptr<T> { return std::make_shared<T>(args...); };
+        this->resize(pool_size_);
+    }
+
+    void
+    SetConstructor(ConstructFuncType func) {
+        this->constructor_ = func;
+        {
+            std::unique_lock<std::mutex> lock(mutex_);
+            while (not pool_.empty()) {
+                pool_.pop_back();
+            }
+        }
+        this->resize(pool_size_);
+    }
+
+    std::shared_ptr<T>
+    GetOne() {
+        std::unique_lock<std::mutex> lock(mutex_);
+        if (pool_.empty()) {
+            return this->constructor_();
+        }
+        std::shared_ptr<T> obj = pool_.back();
+        pool_.pop_back();
+        pool_size_--;
+        obj->Reset();
+        return obj;
+    }
+
+    void
+    ReleaseOne(std::shared_ptr<T>& obj) {
+        std::unique_lock<std::mutex> lock(mutex_);
+        pool_.push_back(obj);
+        pool_size_++;
+    }
+
+    [[nodiscard]] inline uint64_t
+    GetSize() const {
+        return this->pool_size_;
+    }
+
+private:
+    inline void
+    resize(uint64_t size) {
+        std::unique_lock<std::mutex> lock(mutex_);
+        int count = size - pool_.size();
+        while (count > 0) {
+            pool_.emplace_back(this->constructor_());
+            --count;
+        }
+        while (count < 0) {
+            pool_.pop_back();
+            ++count;
+        }
+    }
+
+    std::vector<std::shared_ptr<T>> pool_{};
+    size_t pool_size_{0};
+    ConstructFuncType constructor_{nullptr};
+    std::mutex mutex_;
+};
+
+}  // namespace vsag
diff --git a/src/utils/visited_list.h b/src/utils/visited_list.h
@@ -0,0 +1,80 @@
+
+// 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 <cstring>
+#include <limits>
+
+#include "resource_object.h"
+#include "resource_object_pool.h"
+#include "typing.h"
+#include "vsag/allocator.h"
+
+namespace vsag {
+
+class VisitedList : public ResourceObject {
+public:
+    using VisitedListType = uint16_t;
+
+public:
+    explicit VisitedList(InnerIdType max_size, Allocator* allocator)
+        : max_size_(max_size), allocator_(allocator) {
+        this->list_ = reinterpret_cast<VisitedListType*>(
+            allocator_->Allocate((uint64_t)max_size * sizeof(VisitedListType)));
+        memset(list_, 0, max_size_ * sizeof(VisitedListType));
+        tag_ = 1;
+    }
+
+    ~VisitedList() override {
+        allocator_->Deallocate(list_);
+    }
+
+    inline void
+    Set(const InnerIdType& id) {
+        this->list_[id] = this->tag_;
+    }
+
+    inline bool
+    Get(const InnerIdType& id) {
+        return this->list_[id] == this->tag_;
+    }
+
+    inline void
+    Prefetch(const InnerIdType& id) {
+        return;  // TODO(LHT) implement
+    }
+
+    void
+    Reset() override {
+        if (tag_ == std::numeric_limits<VisitedListType>::max()) {
+            memset(list_, 0, max_size_ * sizeof(VisitedListType));
+            tag_ = 0;
+        }
+        ++tag_;
+    }
+
+private:
+    Allocator* const allocator_{nullptr};
+
+    VisitedListType* list_{nullptr};
+
+    VisitedListType tag_{1};
+
+    const InnerIdType max_size_{0};
+};
+
+using VisitedListPool = ResourceObjectPool<VisitedList>;
+
+}  // namespace vsag