Fix for concurrency bug (#70)

* update * fix after comment * fix after comment * revert fix after comment
spotify · Jul 16, 2024 · 3fc184a · 3fc184a
1 parent 9770a04
commit 3fc184a
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/cpp/TypedIndex.h b/cpp/TypedIndex.h
@@ -99,7 +99,7 @@ class TypedIndex : public Index {
   bool normalize = false;
   bool useOrderPreservingTransform = false;
   int numThreadsDefault;
-  hnswlib::labeltype currentLabel;
+  std::atomic<hnswlib::labeltype> currentLabel;
   std::unique_ptr<hnswlib::HierarchicalNSW<dist_t, data_t>> algorithmImpl;
   std::unique_ptr<hnswlib::Space<dist_t, data_t>> spaceImpl;
   std::unique_ptr<voyager::Metadata::V1> metadata;
@@ -360,7 +360,7 @@ class TypedIndex : public Index {
 
     int start = 0;
     if (!ep_added) {
-      size_t id = ids.size() ? ids.at(0) : (currentLabel);
+      size_t id = ids.size() ? ids.at(0) : (currentLabel.fetch_add(1));
       // TODO(psobot): Should inputVector be on the stack instead?
       std::vector<float> inputVector(actualDimensions);
       std::vector<data_t> convertedVector(actualDimensions);
@@ -402,7 +402,7 @@ class TypedIndex : public Index {
         floatToDataType<data_t, scalefactor>(&inputArray[startIndex],
                                              &convertedArray[startIndex],
                                              actualDimensions);
-        size_t id = ids.size() ? ids.at(row) : (currentLabel + row);
+        size_t id = ids.size() ? ids.at(row) : (currentLabel.fetch_add(1));
         try {
           algorithmImpl->addPoint(convertedArray.data() + startIndex, id);
         } catch (IndexFullError &e) {
@@ -438,7 +438,7 @@ class TypedIndex : public Index {
         normalizeVector<dist_t, data_t, scalefactor>(
             &inputArray[startIndex], &normalizedArray[startIndex],
             actualDimensions);
-        size_t id = ids.size() ? ids.at(row) : (currentLabel + row);
+        size_t id = ids.size() ? ids.at(row) : (currentLabel.fetch_add(1));
 
         try {
           algorithmImpl->addPoint(normalizedArray.data() + startIndex, id);
@@ -460,8 +460,6 @@ class TypedIndex : public Index {
       });
     };
 
-    currentLabel += rows;
-
     return idsToReturn;
   }