8000 Fix memoty leak in HnswHandler. by 19-hanhan · Pull Request #2739 · infiniflow/infinity · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

Fix memoty leak in HnswHandler. #2739

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jul 1, 2025
Merged

Fix memoty leak in HnswHandler. #2739

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
5c56201
Convert HnswHandler::AbstractHnsw from raw pointer to SharedPtr.
19-hanhan Jun 25, 2025
8d6444e
Convert BMPHandler::AbstractBMP from raw pointer to SharedPtr.
19-hanhan Jun 25, 2025
c4a863d
Fix Compile Errors in BMPIndexTest.
19-hanhan Jun 26, 2025
2355967
Add info of INDEX_BUILD_METHOD.
19-hanhan Jun 26, 2025
db1e7bf
Fix BMPAlg Impact of Load & LoadFromPtr in AbstractBMP.
19-hanhan Jun 26, 2025
8a314a8
Fix handler pointer type.
19-hanhan Jun 30, 2025
b21c8ae
Fix destructor of file_worker.
19-hanhan Jul 1, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,10 @@ ADD_DEFINITIONS(-D INFINITY_DEBUG)
endif()

if(NOT INDEX_BUILD_METHOD STREQUAL "Variant")
ADD_DEFINITIONS(-D INDEX_HANDLER)
message(STATUS "Add Environment Variable INDEX_HANDLER")
ADD_DEFINITIONS(-D INDEX_HANDLER)
else()
message(STATUS "NOT USING INDEX_HANDLER")
endif()

# find_package(Boost REQUIRED)
Expand Down
10 changes: 7 additions & 3 deletions src/storage/buffer/file_worker/bmp_index_file_worker.cpp
View file
Open in desktop
8000
Original file line number Diff line numberDiff line change
Expand Up @@ -67,6 +67,10 @@ BMPIndexFileWorker::~BMPIndexFileWorker() {
FreeInMemory();
data_ = nullptr;
}
if (mmap_data_ != nullptr) {
FreeFromMmapImpl();
mmap_data_ = nullptr;
}
}

void BMPIndexFileWorker::AllocateInMemory() {
Expand Down Expand Up @@ -168,9 +172,9 @@ void BMPIndexFileWorker::ReadFromFileImpl(SizeT file_size, bool from_spill) {
using IndexT = std::decay_t<decltype(*index)>;
if constexpr (IndexT::kOwnMem) {
if (from_spill) {
index = new IndexT(IndexT::Load(*file_handle_));
index = IndexT::Load(*file_handle_).release();
} else {
index = new IndexT(IndexT::LoadFromPtr(*file_handle_, file_size));
index = IndexT::LoadFromPtr(*file_handle_, file_size).release();
}
} else {
UnrecoverableError("BMPIndexFileWorker::ReadFromFileImpl: index does not own memory");
Expand Down Expand Up @@ -201,7 +205,7 @@ bool BMPIndexFileWorker::ReadFromMmapImpl(const void *ptr, SizeT size) {
using IndexT = std::decay_t<decltype(*index)>;
if constexpr (!IndexT::kOwnMem) {
const auto *p = static_cast<const char *>(ptr);
index = new IndexT(IndexT::LoadFromPtr(p, size));
index = IndexT::LoadFromPtr(p, size).release();
} else {
UnrecoverableError("BMPIndexFileWorker::ReadFromMmapImpl: index own memory");
}
Expand Down
4 changes: 4 additions & 0 deletions src/storage/buffer/file_worker/hnsw_file_worker.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,10 @@ HnswFileWorker::~HnswFileWorker() {
FreeInMemory();
data_ = nullptr;
}
if (mmap_data_ != nullptr) {
FreeFromMmapImpl();
mmap_data_ = nullptr;
}
}

void HnswFileWorker::AllocateInMemory() {
Expand Down
50 changes: 17 additions & 33 deletions src/storage/knn_index/knn_hnsw/hnsw_handler.cpp
View file
Open in desktop
6D47
Original file line number Diff line number Diff line change
Expand Up @@ -29,15 +29,15 @@ AbstractHnsw InitAbstractIndexT(const IndexHnsw *index_hnsw) {
switch (index_hnsw->metric_type_) {
case MetricType::kMetricL2: {
using HnswIndex = KnnHnsw<PlainL2VecStoreType<DataType, true>, SegmentOffset, OwnMem>;
return static_cast<HnswIndex *>(nullptr);
return UniquePtr<HnswIndex>();
}
case MetricType::kMetricInnerProduct: {
using HnswIndex = KnnHnsw<PlainIPVecStoreType<DataType, true>, SegmentOffset, OwnMem>;
return static_cast<HnswIndex *>(nullptr);
return UniquePtr<HnswIndex>();
}
case MetricType::kMetricCosine: {
using HnswIndex = KnnHnsw<PlainCosVecStoreType<DataType, true>, SegmentOffset, OwnMem>;
return static_cast<HnswIndex *>(nullptr);
return UniquePtr<HnswIndex>();
}
default: {
return nullptr;
Expand All @@ -49,15 +49,15 @@ AbstractHnsw InitAbstractIndexT(const IndexHnsw *index_hnsw) {
switch (index_hnsw->metric_type_) {
case MetricType::kMetricL2: {
using HnswIndex = KnnHnsw<PlainL2VecStoreType<DataType>, SegmentOffset, OwnMem>;
return static_cast<HnswIndex *>(nullptr);
return UniquePtr<HnswIndex>();
}
case MetricType::kMetricInnerProduct: {
using HnswIndex = KnnHnsw<PlainIPVecStoreType<DataType>, SegmentOffset, OwnMem>;
return static_cast<HnswIndex *>(nullptr);
return UniquePtr<HnswIndex>();
}
case MetricType::kMetricCosine: {
using HnswIndex = KnnHnsw<PlainCosVecStoreType<DataType>, SegmentOffset, OwnMem>;
return static_cast<HnswIndex *>(nullptr);
return UniquePtr<HnswIndex>();
}
default: {
return nullptr;
Expand All @@ -71,15 +71,15 @@ AbstractHnsw InitAbstractIndexT(const IndexHnsw *index_hnsw) {
switch (index_hnsw->metric_type_) {
case MetricType::kMetricL2: {
using HnswIndex = KnnHnsw<LVQL2VecStoreType<DataType, i8>, SegmentOffset, OwnMem>;
return static_cast<HnswIndex *>(nullptr);
return UniquePtr<HnswIndex>();
}
case MetricType::kMetricInnerProduct: {
using HnswIndex = KnnHnsw<LVQIPVecStoreType<DataType, i8>, SegmentOffset, OwnMem>;
return static_cast<HnswIndex *>(nullptr);
return UniquePtr<HnswIndex>();
}
case MetricType::kMetricCosine: {
using HnswIndex = KnnHnsw<LVQCosVecStoreType<DataType, i8>, SegmentOffset, OwnMem>;
return static_cast<HnswIndex *>(nullptr);
return UniquePtr<HnswIndex>();
}
default: {
return nullptr;
Expand Down Expand Up @@ -114,17 +114,6 @@ AbstractHnsw InitAbstractIndexT(const IndexBase *index_base, const ColumnDef *co
}
}

HnswHandler::~HnswHandler() {
std::visit(
[&](auto &&index) {
using T = std::decay_t<decltype(index)>;
if constexpr (!std::is_same_v<T, std::nullptr_t>) {
delete index;
}
},
hnsw_);
}

AbstractHnsw HnswHandler::InitAbstractIndex(const IndexBase *index_base, const ColumnDef *column_def, bool own_mem) {
if (own_mem) {
return InitAbstractIndexT<true>(index_base, column_def);
Expand Down Expand Up @@ -152,7 +141,7 @@ HnswHandler::HnswHandler(const IndexBase *index_base, const ColumnDef *column_de
if constexpr (!std::is_same_v<T, std::nullptr_t>) {
using IndexT = std::decay_t<decltype(*index)>;
if constexpr (IndexT::kOwnMem) {
index = IndexT::Make(chunk_size, max_chunk_num, dim, M, ef_construction).release();
index = IndexT::Make(chunk_size, max_chunk_num, dim, M, ef_construction);
} else {
UnrecoverableError("HnswHandler::HnswHandler: index does not own memory");
}
Expand Down Expand Up @@ -237,7 +226,7 @@ SizeT HnswHandler::GetSizeInBytes() const {
if constexpr (std::is_same_v<T, std::nullptr_t>) {
return SizeT(0);
} else {
using IndexT = typename std::remove_pointer_t<T>;
using IndexT = std::decay_t<decltype(*index)>;
if constexpr (IndexT::kOwnMem) {
return index->GetSizeInBytes();
} else {
Expand Down Expand Up @@ -308,8 +297,7 @@ void HnswHandler::Load(LocalFileHandle &file_handle) {
} else {
using IndexT = std::decay_t<decltype(*index)>;
if constexpr (IndexT::kOwnMem) {
delete index;
index = IndexT::Load(file_handle).release();
index = IndexT::Load(file_handle);
} else {
UnrecoverableError("Invalid index type.");
}
Expand All @@ -327,8 +315,7 @@ void HnswHandler::LoadFromPtr(LocalFileHandle &file_handle, SizeT file_size) {
} else {
using IndexT = std::decay_t<decltype(*index)>;
if constexpr (IndexT::kOwnMem) {
delete index;
index = IndexT::LoadFromPtr(file_handle, file_size).release();
index = IndexT::LoadFromPtr(file_handle, file_size);
} else {
UnrecoverableError("Invalid index type.");
}
Expand All @@ -348,8 +335,7 @@ void HnswHandler::LoadFromPtr(const char *ptr, SizeT size) {
if constexpr (IndexT::kOwnMem) {
UnrecoverableError("Invalid index type.");
} else {
delete index;
index = IndexT::LoadFromPtr(ptr, size).release();
index = IndexT::LoadFromPtr(ptr, size);
}
}
},
Expand Down Expand Up @@ -399,15 +385,13 @@ void HnswHandler::CompressToLVQ() {
if constexpr (std::is_same_v<T, std::nullptr_t>) {
UnrecoverableError("Invalid index type.");
} else {
using IndexT = typename std::remove_pointer_t<T>;
using IndexT = std::decay_t<decltype(*index)>;
if constexpr (IndexT::kOwnMem) {
using HnswIndexDataType = typename std::remove_pointer_t<T>::DataType;
using HnswIndexDataType = IndexT::DataType;
if constexpr (IsAnyOf<HnswIndexDataType, i8, u8>) {
UnrecoverableError("Invalid index type.");
} else {
auto *p = std::move(*index).CompressToLVQ().release();
delete index;
hnsw_ = p;
hnsw_ = std::move(*index).CompressToLVQ();
}
} else {
UnrecoverableError("Invalid index type.");
Expand Down
89 changes: 44 additions & 45 deletions src/storage/knn_index/knn_hnsw/hnsw_handler.cppm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,56 +40,55 @@ class ColumnVector;
class BufferObj;
class LocalFileHandle;

using AbstractHnsw = std::variant<KnnHnsw<PlainCosVecStoreType<float>, SegmentOffset> *,
KnnHnsw<PlainIPVecStoreType<float>, SegmentOffset> *,
KnnHnsw<PlainL2VecStoreType<float>, SegmentOffset> *,
KnnHnsw<PlainCosVecStoreType<u8>, SegmentOffset> *,
KnnHnsw<PlainIPVecStoreType<u8>, SegmentOffset> *,
KnnHnsw<PlainL2VecStoreType<u8>, SegmentOffset> *,
KnnHnsw<PlainCosVecStoreType<i8>, SegmentOffset> *,
KnnHnsw<PlainIPVecStoreType<i8>, SegmentOffset> *,
KnnHnsw<PlainL2VecStoreType<i8>, SegmentOffset> *,
KnnHnsw<LVQCosVecStoreType<float, i8>, SegmentOffset> *,
KnnHnsw<LVQIPVecStoreType<float, i8>, SegmentOffset> *,
KnnHnsw<LVQL2VecStoreType<float, i8>, SegmentOffset> *,
KnnHnsw<PlainCosVecStoreType<float, true>, SegmentOffset> *,
KnnHnsw<PlainIPVecStoreType<float, true>, SegmentOffset> *,
KnnHnsw<PlainL2VecStoreType<float, true>, SegmentOffset> *,
KnnHnsw<PlainCosVecStoreType<u8, true>, SegmentOffset> *,
KnnHnsw<PlainIPVecStoreType<u8, true>, SegmentOffset> *,
KnnHnsw<PlainL2VecStoreType<u8, true>, SegmentOffset> *,
KnnHnsw<PlainCosVecStoreType<i8, true>, SegmentOffset> *,
KnnHnsw<PlainIPVecStoreType<i8, true>, SegmentOffset> *,
KnnHnsw<PlainL2VecStoreType<i8, true>, SegmentOffset> *,

KnnHnsw<PlainCosVecStoreType<float>, SegmentOffset, false> *,
KnnHnsw<PlainIPVecStoreType<float>, SegmentOffset, false> *,
KnnHnsw<PlainL2VecStoreType<float>, SegmentOffset, false> *,
KnnHnsw<PlainCosVecStoreType<u8>, SegmentOffset, false> *,
KnnHnsw<PlainIPVecStoreType<u8>, SegmentOffset, false> *,
KnnHnsw<PlainL2VecStoreType<u8>, SegmentOffset, false> *,
KnnHnsw<PlainCosVecStoreType<i8>, SegmentOffset, false> *,
KnnHnsw<PlainIPVecStoreType<i8>, SegmentOffset, false> *,
KnnHnsw<PlainL2VecStoreType<i8>, SegmentOffset, false> *,
KnnHnsw<LVQCosVecStoreType<float, i8>, SegmentOffset, false> *,
KnnHnsw<LVQIPVecStoreType<float, i8>, SegmentOffset, false> *,
KnnHnsw<LVQL2VecStoreType<float, i8>, SegmentOffset, false> *,
KnnHnsw<PlainCosVecStoreType<float, true>, SegmentOffset, false> *,
KnnHnsw<PlainIPVecStoreType<float, true>, SegmentOffset, false> *,
KnnHnsw<PlainL2VecStoreType<float, true>, SegmentOffset, false> *,
KnnHnsw<PlainCosVecStoreType<u8, true>, SegmentOffset, false> *,
KnnHnsw<PlainIPVecStoreType<u8, true>, SegmentOffset, false> *,
KnnHnsw<PlainL2VecStoreType<u8, true>, SegmentOffset, false> *,
KnnHnsw<PlainCosVecStoreType<i8, true>, SegmentOffset, false> *,
KnnHnsw<PlainIPVecStoreType<i8, true>, SegmentOffset, false> *,
KnnHnsw<PlainL2VecStoreType<i8, true>, SegmentOffset, false> *,
using AbstractHnsw = std::variant<UniquePtr<KnnHnsw<PlainCosVecStoreType<float>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<float>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<float>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainCosVecStoreType<u8>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<u8>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<u8>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainCosVecStoreType<i8>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<i8>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<i8>, SegmentOffset>>,
UniquePtr<KnnHnsw<LVQCosVecStoreType<float, i8>, SegmentOffset>>,
UniquePtr<KnnHnsw<LVQIPVecStoreType<float, i8>, SegmentOffset>>,
UniquePtr<KnnHnsw<LVQL2VecStoreType<float, i8>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainCosVecStoreType<float, true>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<float, true>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<float, true>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainCosVecStoreType<u8, true>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<u8, true>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<u8, true>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainCosVecStoreType<i8, true>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<i8, true>, SegmentOffset>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<i8, true>, SegmentOffset>>,

UniquePtr<KnnHnsw<PlainCosVecStoreType<float>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<float>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<float>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainCosVecStoreType<u8>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<u8>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<u8>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainCosVecStoreType<i8>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<i8>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<i8>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<LVQCosVecStoreType<float, i8>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<LVQIPVecStoreType<float, i8>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<LVQL2VecStoreType<float, i8>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainCosVecStoreType<float, true>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<float, true>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<float, true>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainCosVecStoreType<u8, true>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<u8, true>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<u8, true>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainCosVecStoreType<i8, true>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainIPVecStoreType<i8, true>, SegmentOffset, false>>,
UniquePtr<KnnHnsw<PlainL2VecStoreType<i8, true>, SegmentOffset, false>>,
std::nullptr_t>;

export struct HnswHandler {
public:
HnswHandler() : hnsw_(nullptr) {}
HnswHandler(const void *ptr) : hnsw_(*reinterpret_cast<const AbstractHnsw *>(ptr)) {}
virtual ~HnswHandler();
virtual ~HnswHandler() {}
HnswHandler(const HnswHandler &) = delete;
HnswHandler &operator=(const HnswHandler &) = delete;

Expand Down
Loading
0