Add lock to avoid segfault while resizing (#61)

* Add lock avoid segfault while resizing * clang-fmt * try read-write lock * add rw lock * acquire shared lock around entire searchKnn method * clang-fmt * target macOS 10.13 in python build * use shared_ptr/unique_ptr to take the lock * lock on addPoint * revert using loop to resize index * Ensure index resizes still happen correctly if adding in multiple threads. * formatting * fix merge * fmt * more formatting --------- Co-authored-by: Peter Sobot <[email protected]>
spotify · Mar 22, 2024 · cc02679 · cc02679
1 parent 63b8d60
commit cc02679
Show file tree

Hide file tree

Showing 38 changed files with 258 additions and 80 deletions.
diff --git a/cpp/TypedIndex.h b/cpp/TypedIndex.h
@@ -346,8 +346,13 @@ class TypedIndex : public Index {
 
     // TODO: Should we always double the number of elements instead? Maybe use
     // an adaptive algorithm to minimize both reallocations and memory usage?
-    if (getNumElements() + rows > getMaxElements()) {
-      resizeIndex(getNumElements() + rows);
+    while (getNumElements() + rows > getMaxElements()) {
+      try {
+        resizeIndex(getNumElements() + rows);
+      } catch (IndexCannotBeShrunkError &e) {
+        // Retry with a larger size; some other thread may have resized
+        // behind our back.
+      }
     }
 
     int actualDimensions =
@@ -398,7 +403,22 @@ class TypedIndex : public Index {
                                              &convertedArray[startIndex],
                                              actualDimensions);
         size_t id = ids.size() ? ids.at(row) : (currentLabel + row);
-        algorithmImpl->addPoint(convertedArray.data() + startIndex, id);
+        try {
+          algorithmImpl->addPoint(convertedArray.data() + startIndex, id);
+        } catch (IndexFullError &e) {
+          // Resize the index and try again:
+          while (getNumElements() + rows > getMaxElements()) {
+            try {
+              // NOTE: This will resize the index to be at least as large as
+              // the number of elements we're trying to add, but may
+              // allocate more space than necessary.
+              resizeIndex(getNumElements() + rows);
+            } catch (IndexCannotBeShrunkError &e) {
+              // Retry with a larger size; some other thread may have resized
+              // behind our back.
+            }
+          }
+        }
         idsToReturn[row] = id;
       });
     } else {
@@ -419,7 +439,23 @@ class TypedIndex : public Index {
             &inputArray[startIndex], &normalizedArray[startIndex],
             actualDimensions);
         size_t id = ids.size() ? ids.at(row) : (currentLabel + row);
-        algorithmImpl->addPoint(normalizedArray.data() + startIndex, id);
+
+        try {
+          algorithmImpl->addPoint(normalizedArray.data() + startIndex, id);
+        } catch (IndexFullError &e) {
+          // Resize the index and try again:
+          while (getNumElements() + rows > getMaxElements()) {
+            try {
+              // NOTE: This will resize the index to be at least as large as
+              // the number of elements we're trying to add, but may
+              // allocate more space than necessary.
+              resizeIndex(getNumElements() + rows);
+            } catch (IndexCannotBeShrunkError &e) {
+              // Retry with a larger size; some other thread may have resized
+              // behind our back.
+            }
+          }
+        }
         idsToReturn[row] = id;
       });
     };

diff --git a/cpp/hnswalg.h b/cpp/hnswalg.h
@@ -33,10 +33,22 @@
 #include <fstream>
 #include <list>
 #include <random>
+#include <shared_mutex>
 #include <stdlib.h>
 #include <unordered_map>
 #include <unordered_set>
 
+class IndexCannotBeShrunkError : public std::runtime_error {
+public:
+  IndexCannotBeShrunkError(const std::string &what)
+      : std::runtime_error(what) {}
+};
+
+class IndexFullError : public std::runtime_error {
+public:
+  IndexFullError(const std::string &what) : std::runtime_error(what) {}
+};
+
 namespace hnswlib {
 typedef unsigned int tableint;
 typedef unsigned int linklistsizeint;
@@ -135,6 +147,7 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t, data_t> {
   double mult_, revSize_;
   int maxlevel_;
 
+  std::shared_mutex resizeLock;
   VisitedListPool *visited_list_pool_;
   std::mutex cur_element_count_guard_;
 
@@ -642,10 +655,13 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t, data_t> {
     if (search_only_)
       throw std::runtime_error(
           "resizeIndex is not supported in search only mode");
+    std::unique_lock<std::shared_mutex> lock(resizeLock);
 
     if (new_max_elements < cur_element_count)
-      throw std::runtime_error("Cannot resize, max element is less than the "
-                               "current number of elements");
+      throw IndexCannotBeShrunkError(
+          "Cannot resize to " + std::to_string(new_max_elements) +
+          " elements, as this index already contains " +
+          std::to_string(cur_element_count) + " elements.");
 
     delete visited_list_pool_;
     visited_list_pool_ = new VisitedListPool(1, new_max_elements);
@@ -1262,7 +1278,7 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t, data_t> {
   };
 
   tableint addPoint(const data_t *data_point, labeltype label, int level) {
-
+    std::shared_lock<std::shared_mutex> lock(resizeLock);
     tableint cur_c = 0;
     {
       // Checking if the element with the same label already exists
@@ -1285,8 +1301,13 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t, data_t> {
       }
 
       if (cur_element_count >= max_elements_) {
-        throw std::runtime_error(
-            "The number of elements exceeds the specified limit");
+        throw IndexFullError(
+            "Cannot insert elements; this index already contains " +
+            std::to_string(cur_element_count) +
+            " elements, and its maximum size is " +
+            std::to_string(max_elements_) +
+            ". Call resizeIndex first to increase the maximum size of the "
+            "index.");
       };
 
       cur_c = cur_element_count;
@@ -1398,7 +1419,8 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t, data_t> {
 
   std::priority_queue<std::pair<dist_t, labeltype>>
   searchKnn(const data_t *query_data, size_t k, VisitedList *vl = nullptr,
-            long queryEf = -1) const {
+            long queryEf = -1) {
+    std::shared_lock<std::shared_mutex> lock(resizeLock);
     std::priority_queue<std::pair<dist_t, labeltype>> result;
     if (cur_element_count == 0)
       return result;

diff --git a/cpp/hnswlib.h b/cpp/hnswlib.h
@@ -77,11 +77,11 @@ template <typename dist_t, typename data_t = dist_t> class AlgorithmInterface {
   virtual void addPoint(const data_t *datapoint, labeltype label) = 0;
   virtual std::priority_queue<std::pair<dist_t, labeltype>>
   searchKnn(const data_t *, size_t, VisitedList *a = nullptr,
-            long queryEf = -1) const = 0;
+            long queryEf = -1) = 0;
 
   // Return k nearest neighbor in the order of closer fist
   virtual std::vector<std::pair<dist_t, labeltype>>
-  searchKnnCloserFirst(const data_t *query_data, size_t k) const;
+  searchKnnCloserFirst(const data_t *query_data, size_t k);
 
   virtual void saveIndex(const std::string &location) = 0;
   virtual ~AlgorithmInterface() {}
@@ -90,7 +90,7 @@ template <typename dist_t, typename data_t = dist_t> class AlgorithmInterface {
 template <typename dist_t, typename data_t>
 std::vector<std::pair<dist_t, labeltype>>
 AlgorithmInterface<dist_t, data_t>::searchKnnCloserFirst(
-    const data_t *query_data, size_t k) const {
+    const data_t *query_data, size_t k) {
   std::vector<std::pair<dist_t, labeltype>> result;
 
   // here searchKnn returns the result in the order of further first

diff --git a/docs/java/allclasses-index.html b/docs/java/allclasses-index.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>All Classes and Interfaces (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="class index">
 <meta name="generator" content="javadoc/AllClassesIndexWriter">
 <link rel="stylesheet" type="text/css" href="stylesheet.css" title="Style">

diff --git a/docs/java/allpackages-index.html b/docs/java/allpackages-index.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>All Packages (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="package index">
 <meta name="generator" content="javadoc/AllPackagesIndexWriter">
 <link rel="stylesheet" type="text/css" href="stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/Index.QueryResults.html b/docs/java/com/spotify/voyager/jni/Index.QueryResults.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>Index.QueryResults (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="declaration: package: com.spotify.voyager.jni, class: Index, class: QueryResults">
 <meta name="generator" content="javadoc/ClassWriterImpl">
 <link rel="stylesheet" type="text/css" href="../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/Index.SpaceType.html b/docs/java/com/spotify/voyager/jni/Index.SpaceType.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>Index.SpaceType (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="declaration: package: com.spotify.voyager.jni, class: Index, enum: SpaceType">
 <meta name="generator" content="javadoc/ClassWriterImpl">
 <link rel="stylesheet" type="text/css" href="../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/Index.StorageDataType.html b/docs/java/com/spotify/voyager/jni/Index.StorageDataType.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>Index.StorageDataType (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="declaration: package: com.spotify.voyager.jni, class: Index, enum: StorageDataType">
 <meta name="generator" content="javadoc/ClassWriterImpl">
 <link rel="stylesheet" type="text/css" href="../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/Index.html b/docs/java/com/spotify/voyager/jni/Index.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>Index (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="declaration: package: com.spotify.voyager.jni, class: Index">
 <meta name="generator" content="javadoc/ClassWriterImpl">
 <link rel="stylesheet" type="text/css" href="../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/StringIndex.QueryResults.html b/docs/java/com/spotify/voyager/jni/StringIndex.QueryResults.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>StringIndex.QueryResults (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="declaration: package: com.spotify.voyager.jni, class: StringIndex, class: QueryResults">
 <meta name="generator" content="javadoc/ClassWriterImpl">
 <link rel="stylesheet" type="text/css" href="../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/StringIndex.html b/docs/java/com/spotify/voyager/jni/StringIndex.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>StringIndex (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="declaration: package: com.spotify.voyager.jni, class: StringIndex">
 <meta name="generator" content="javadoc/ClassWriterImpl">
 <link rel="stylesheet" type="text/css" href="../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/class-use/Index.QueryResults.html b/docs/java/com/spotify/voyager/jni/class-use/Index.QueryResults.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>Uses of Class com.spotify.voyager.jni.Index.QueryResults (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="use: package: com.spotify.voyager.jni, class: Index, class: QueryResults">
 <meta name="generator" content="javadoc/ClassUseWriter">
 <link rel="stylesheet" type="text/css" href="../../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/class-use/Index.SpaceType.html b/docs/java/com/spotify/voyager/jni/class-use/Index.SpaceType.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>Uses of Enum com.spotify.voyager.jni.Index.SpaceType (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="use: package: com.spotify.voyager.jni, class: Index, enum: SpaceType">
 <meta name="generator" content="javadoc/ClassUseWriter">
 <link rel="stylesheet" type="text/css" href="../../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/class-use/Index.StorageDataType.html b/docs/java/com/spotify/voyager/jni/class-use/Index.StorageDataType.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>Uses of Enum com.spotify.voyager.jni.Index.StorageDataType (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="use: package: com.spotify.voyager.jni, class: Index, enum: StorageDataType">
 <meta name="generator" content="javadoc/ClassUseWriter">
 <link rel="stylesheet" type="text/css" href="../../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/class-use/Index.html b/docs/java/com/spotify/voyager/jni/class-use/Index.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>Uses of Class com.spotify.voyager.jni.Index (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="use: package: com.spotify.voyager.jni, class: Index">
 <meta name="generator" content="javadoc/ClassUseWriter">
 <link rel="stylesheet" type="text/css" href="../../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/class-use/StringIndex.QueryResults.html b/docs/java/com/spotify/voyager/jni/class-use/StringIndex.QueryResults.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>Uses of Class com.spotify.voyager.jni.StringIndex.QueryResults (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="use: package: com.spotify.voyager.jni, class: StringIndex, class: QueryResults">
 <meta name="generator" content="javadoc/ClassUseWriter">
 <link rel="stylesheet" type="text/css" href="../../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/class-use/StringIndex.html b/docs/java/com/spotify/voyager/jni/class-use/StringIndex.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>Uses of Class com.spotify.voyager.jni.StringIndex (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="use: package: com.spotify.voyager.jni, class: StringIndex">
 <meta name="generator" content="javadoc/ClassUseWriter">
 <link rel="stylesheet" type="text/css" href="../../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/package-summary.html b/docs/java/com/spotify/voyager/jni/package-summary.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>com.spotify.voyager.jni (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="declaration: package: com.spotify.voyager.jni">
 <meta name="generator" content="javadoc/PackageWriterImpl">
 <link rel="stylesheet" type="text/css" href="../../../../stylesheet.css" title="Style">

diff --git a/docs/java/com/spotify/voyager/jni/package-tree.html b/docs/java/com/spotify/voyager/jni/package-tree.html
@@ -1,11 +1,11 @@
 <!DOCTYPE HTML>
 <html lang="en">
 <head>
-<!-- Generated by javadoc (17) on Wed Mar 20 12:39:31 EDT 2024 -->
+<!-- Generated by javadoc (17) on Thu Mar 21 17:31:39 EDT 2024 -->
 <title>com.spotify.voyager.jni Class Hierarchy (voyager 2.0.5 API)</title>
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
-<meta name="dc.created" content="2024-03-20">
+<meta name="dc.created" content="2024-03-21">
 <meta name="description" content="tree: package: com.spotify.voyager.jni">
 <meta name="generator" content="javadoc/PackageTreeWriter">
 <link rel="stylesheet" type="text/css" href="../../../../stylesheet.css" title="Style">