Fix thread safety in hnsw (#575)

Signed-off-by: zh Wang <[email protected]> Signed-off-by: zh Wang <[email protected]>
milvus-io · Dec 1, 2022 · 998786c · 998786c
1 parent 2ac540e
commit 998786c
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/thirdparty/hnswlib/hnswlib/visited_list_pool.h b/thirdparty/hnswlib/hnswlib/visited_list_pool.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <algorithm>
+#include <mutex>
 #include <thread>
 #include <unordered_map>
 #include <vector>
@@ -16,6 +17,7 @@ namespace hnswlib {
 class VisitedListPool {
     int numelements;
     std::unordered_map<std::thread::id, std::vector<bool>> map;
+    std::mutex mtx;
 
  public:
     VisitedListPool(int numelements1) {
@@ -24,7 +26,9 @@ class VisitedListPool {
 
     std::vector<bool>&
     getFreeVisitedList() {
+        std::unique_lock lk(mtx);
         auto& res = map[std::this_thread::get_id()];
+        lk.unlock();
         if (res.size() != numelements) {
             res.assign(numelements, false);
         } else {