Skip to content

Commit b11eea9

Browse files
LHT129LHT129
committed
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]>
1 parent 0f7e95c commit b11eea9

File tree

5 files changed

+346
-1
lines changed

5 files changed

+346
-1
lines changed

‎src/CMakeLists.txt

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,11 @@ list (FILTER CPP_INDEX_SRCS EXCLUDE REGEX "_test.cpp")
1515
list (FILTER CPP_QUANTIZATION_SRCS EXCLUDE REGEX "_test.cpp")
1616
list (FILTER CPP_DATA_CELL_SRCS EXCLUDE REGEX "_test.cpp")
1717

18+
file (GLOB CPP_UTILS_SRCS "*.cpp")
19+
list (FILTER CPP_UTILS_SRCS EXCLUDE REGEX "_test.cpp")
20+
1821
set (VSAG_SRCS ${CPP_SRCS} ${CPP_FACTORY_SRCS} ${CPP_INDEX_SRCS}
19-
${CPP_CONJUGATE_GRAPH_SRCS} ${CPP_HNSWLIB_SRCS} ${CPP_QUANTIZATION_SRCS} ${CPP_DATA_CELL_SRCS})
22+
${CPP_CONJUGATE_GRAPH_SRCS} ${CPP_HNSWLIB_SRCS} ${CPP_QUANTIZATION_SRCS} ${CPP_DATA_CELL_SRCS} ${CPP_UTILS_SRCS})
2023
add_library (vsag SHARED ${VSAG_SRCS})
2124
add_library (vsag_static STATIC ${VSAG_SRCS})
2225

‎src/utils/resource_object.h

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
2+
// Copyright 2024-present the vsag project
3+
//
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
//
8+
// http://www.apache.org/licenses/LICENSE-2.0
9+
//
10+
// Unless required by applicable law or agreed to in writing, software
11+
// distributed under the License is distributed on an "AS IS" BASIS,
12+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
// See the License for the specific language governing permissions and
14+
// limitations under the License.
15+
16+
#pragma once
17+
18+
namespace vsag {
19+
20+
class ResourceObject {
21+
public:
22+
ResourceObject() = default;
23+
24+
virtual ~ResourceObject() = default;
25+
26+
virtual void
27+
Reset() = 0;
28+
};
29+
30+
} // namespace vsag

‎src/utils/resource_object_pool.h

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
2+
// Copyright 2024-present the vsag project
3+
//
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
//
8+
// http://www.apache.org/licenses/LICENSE-2.0
9+
//
10+
// Unless required by applicable law or agreed to in writing, software
11+
// distributed under the License is distributed on an "AS IS" BASIS,
12+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
// See the License for the specific language governing permissions and
14+
// limitations under the License.
15+
16+
#pragma once
17+
18+
#include <cstdint>
19+
#include <functional>
20+
#include <memory>
21+
#include <mutex>
22+
#include <type_traits>
23+
24+
#include "resource_object.h"
25+
namespace vsag {
26+
27+
template <typename T,
28+
typename = typename std::enable_if<std::is_base_of<ResourceObject, T>::value>::type>
29+
class ResourceObjectPool {
30+
public:
31+
using ConstructFuncType = std::function<std::shared_ptr<T>()>;
32+
33+
public:
34+
template <typename... Args>
35+
explicit ResourceObjectPool(uint64_t init_size, Args... args) {
36+
this->pool_size_ = init_size;
37+
this->constructor_ = [=]() -> std::shared_ptr<T> { return std::make_shared<T>(args...); };
38+
this->resize(pool_size_);
39+
}
40+
41+
void
42+
SetConstructor(ConstructFuncType func) {
43+
this->constructor_ = func;
44+
{
45+
std::unique_lock<std::mutex> lock(mutex_);
46+
while (not pool_.empty()) {
47+
pool_.pop_back();
48+
}
49+
}
50+
this->resize(pool_size_);
51+
}
52+
53+
std::shared_ptr<T>
54+
TakeOne() {
55+
std::unique_lock<std::mutex> lock(mutex_);
56+
if (pool_.empty()) {
57+
return this->constructor_();
58+
}
59+
std::shared_ptr<T> obj = pool_.back();
60+
pool_.pop_back();
61+
pool_size_--;
62+
obj->Reset();
63+
return obj;
64+
}
65+
66+
void
67+
ReturnOne(std::shared_ptr<T>& obj) {
68+
std::unique_lock<std::mutex> lock(mutex_);
69+
pool_.push_back(obj);
70+
pool_size_++;
71+
}
72+
73+
[[nodiscard]] inline uint64_t
74+
GetSize() const {
75+
return this->pool_size_;
76+
}
77+
78+
private:
79+
inline void
80+
resize(uint64_t size) {
81+
std::unique_lock<std::mutex> lock(mutex_);
82+
int count = size - pool_.size();
83+
while (count > 0) {
84+
pool_.emplace_back(this->constructor_());
85+
--count;
86+
}
87+
while (count < 0) {
88+
pool_.pop_back();
89+
++count;
90+
}
91+
}
92+
93+
std::vector<std::shared_ptr<T>> pool_{};
94+
size_t pool_size_{0};
95+
ConstructFuncType constructor_{nullptr};
96+
std::mutex mutex_;
97+
};
98+
99+
} // namespace vsag

‎src/utils/visited_list.h

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
2+
// Copyright 2024-present the vsag project
3+
//
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
//
8+
// http://www.apache.org/licenses/LICENSE-2.0
9+
//
10+
// Unless required by applicable law or agreed to in writing, software
11+
// distributed under the License is distributed on an "AS IS" BASIS,
12+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
// See the License for the specific language governing permissions and
14+
// limitations under the License.
15+
16+
#pragma once
17+
#include <cstring>
18+
#include <limits>
19+
20+
#include "resource_object.h"
21+
#include "resource_object_pool.h"
22+
#include "typing.h"
23+
#include "vsag/allocator.h"
24+
25+
namespace vsag {
26+
27+
class VisitedList : public ResourceObject {
28+
public:
29+
using VisitedListType = uint16_t;
30+
31+
public:
32+
explicit VisitedList(InnerIdType max_size, Allocator* allocator)
33+
: max_size_(max_size), allocator_(allocator) {
34+
this->list_ = reinterpret_cast<VisitedListType*>(
35+
allocator_->Allocate((uint64_t)max_size * sizeof(VisitedListType)));
36+
memset(list_, 0, max_size_ * sizeof(VisitedListType));
37+
tag_ = 1;
38+
}
39+
40+
~VisitedList() override {
41+
allocator_->Deallocate(list_);
42+
}
43+
44+
inline void
45+
Set(const InnerIdType& id) {
46+
this->list_[id] = this->tag_;
47+
}
48+
49+
inline bool
50+
Get(const InnerIdType& id) {
51+
return this->list_[id] == this->tag_;
52+
}
53+
54+
inline void
55+
Prefetch(const InnerIdType& id) {
56+
return; // TODO(LHT) implement
57+
}
58+
59+
void
60+
Reset() override {
61+
if (tag_ == std::numeric_limits<VisitedListType>::max()) {
62+
memset(list_, 0, max_size_ * sizeof(VisitedListType));
63+
tag_ = 0;
64+
}
65+
++tag_;
66+
}
67+
68+
private:
69+
Allocator* const allocator_{nullptr};
70+
71+
VisitedListType* list_{nullptr};
72+
73+
VisitedListType tag_{1};
74+
75+
const InnerIdType max_size_{0};
76+
};
77+
78+
using VisitedListPool = ResourceObjectPool<VisitedList>;
79+
80+
} // namespace vsag

‎src/utils/visited_list_test.cpp

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
2+
// Copyright 2024-present the vsag project
3+
//
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
//
8+
// http://www.apache.org/licenses/LICENSE-2.0
9+
//
10+
// Unless required by applicable law or agreed to in writing, software
11+
// distributed under the License is distributed on an "AS IS" BASIS,
12+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
// See the License for the specific language governing permissions and
14+
// limitations under the License.
15+
16+
#include "visited_list.h"
17+
18+
#include <thread>
19+
20+
#include "catch2/catch_test_macros.hpp"
21+
#include "default_allocator.h"
22+
using namespace vsag;
23+
24+
TEST_CASE("test visited_list basic", "[ut][visited_list]") {
25+
auto allocator = std::make_shared<DefaultAllocator>();
26+
auto size = 10000;
27+
auto vl_ptr = std::make_shared<VisitedList>(size, allocator.get());
28+
29+
SECTION("test set & get normal") {
30+
int count = 500;
31+
std::unordered_set<InnerIdType> ids;
32+
for (int i = 0; i < count; ++i) {
33+
auto id = random() % size;
34+
ids.insert(id);
35+
vl_ptr->Set(id);
36+
}
37+
for (auto& id : ids) {
38+
REQUIRE(vl_ptr->Get(id));
39+
}
40+
41+
for (int i = 0; i < size; ++i) {
42+
if (ids.count(i) == 0) {
43+
REQUIRE(not vl_ptr->Get(i));
44+
}
45+
}
46+
}
47+
48+
SECTION("test reset") {
49+
int count = 500;
50+
std::unordered_set<InnerIdType> ids;
51+
for (int i = 0; i < count; ++i) {
52+
auto id = random() % size;
53+
ids.insert(id);
54+
vl_ptr->Set(id);
55+
}
56+
vl_ptr->Reset();
57+
for (auto& id : ids) {
58+
REQUIRE(not vl_ptr->Get(id));
59+
}
60+
}
61+
}
62+
63+
TEST_CASE("test visited_list_pool basic", "[ut][visited_list_pool]") {
64+
auto allocator = std::make_shared<DefaultAllocator>();
65+
auto init_size = 10;
66+
auto vl_size = 1000;
67+
auto pool = std::make_shared<VisitedListPool>(init_size, vl_size, allocator.get());
68+
69+
auto TestVL = [&](std::shared_ptr<VisitedList>& vl_ptr) {
70+
int count = 500;
71+
std::unordered_set<InnerIdType> ids;
72+
for (int i = 0; i < count; ++i) {
73+
auto id = random() % vl_size;
74+
ids.insert(id);
75+
vl_ptr->Set(id);
76+
}
77+
for (auto& id : ids) {
78+
REQUIRE(vl_ptr->Get(id) == true);
79+
}
80+
81+
for (InnerIdType i = 0; i < vl_size; ++i) {
82+
if (ids.count(i) == 0) {
83+
REQUIRE(vl_ptr->Get(i) == false);
84+
}
85+
}
86+
};
87+
88+
SECTION("test basic") {
89+
std::vector<std::shared_ptr<VisitedList>> lists;
90+
REQUIRE(pool->GetSize() == init_size);
91+
lists.reserve(init_size * 2);
92+
for (auto i = 0; i < init_size * 2; ++i) {
93+
lists.emplace_back(pool->TakeOne());
94+
}
95+
REQUIRE(pool->GetSize() == 0);
96+
for (auto& ptr : lists) {
97+
pool->ReturnOne(ptr);
98+
}
99+
REQUIRE(pool->GetSize() == init_size * 2);
100+
101+
auto ptr = pool->TakeOne();
102+
REQUIRE(pool->GetSize() == init_size * 2 - 1);
103+
TestVL(ptr);
104+
}
105+
106+
SECTION("test concurrency") {
107+
auto func = [&]() {
108+
int count = 10;
109+
int max_operators = 20;
110+
std::vector<std::shared_ptr<VisitedList>> results;
111+
for (int i = 0; i < count; ++i) {
112+
auto opt = random() % max_operators + 1;
113+
for (auto j = 0; j < opt; ++j) {
114+
results.emplace_back(pool->TakeOne());
115+
}
116+
auto test = results[random() % opt];
117+
TestVL(test);
118+
for (auto& result : results) {
119+
pool->ReturnOne(result);
120+
}
121+
results.clear();
122+
}
123+
};
124+
std::vector<std::shared_ptr<std::thread>> ths;
125+
auto thread_count = 2;
126+
for (auto i = 0; i < thread_count; ++i) {
127+
ths.emplace_back((std::make_shared<std::thread>(func)));
128+
}
129+
for (auto& thread : ths) {
130+
thread->join();
131+
}
132+
}
133+
}

0 commit comments

Comments
 (0)