[MOD-10559] Decouple the shrinking and growing logic of large containers in Flat and HNSW

src/VecSim/algorithms/brute_force/brute_force.h

@@ -89,20 +89,45 @@ class BruteForceIndex : public VecSimIndexAbstract<DataType, DistType> {
     // Private internal function that implements generic single vector deletion.
     virtual void removeVector(idType id);
 
-    void growByBlock() {
-        idToLabelMapping.resize(idToLabelMapping.size() + this->blockSize);
+    void resizeIndexCommon(size_t new_max_elements) {
+        assert(new_max_elements % this->blockSize == 0 &&
+               "new_max_elements must be a multiple of blockSize");
+        this->log(VecSimCommonStrings::LOG_VERBOSE_STRING, "Resizing FLAT index from %zu to %zu",
+                  idToLabelMapping.capacity(), new_max_elements);
+        assert(idToLabelMapping.capacity() == idToLabelMapping.size());
+        idToLabelMapping.resize(new_max_elements);
         idToLabelMapping.shrink_to_fit();
-        resizeLabelLookup(idToLabelMapping.size());
+        assert(idToLabelMapping.capacity() == idToLabelMapping.size());
+        resizeLabelLookup(new_max_elements);
     }
 
-    void shrinkByBlock() {
-        assert(indexCapacity() > 0); // should not be called when index is empty
+    void growByBlock() {
+        assert(indexCapacity() == idToLabelMapping.capacity());
+        assert(indexCapacity() % this->blockSize == 0);
+        assert(indexCapacity() == indexSize());
+        assert((dynamic_cast<DataBlocksContainer *>(this->vectors)->numBlocks() ==
+                (indexSize()) / this->blockSize));
 
-        // remove a block size of labels.
-        assert(idToLabelMapping.size() >= this->blockSize);
-        idToLabelMapping.resize(idToLabelMapping.size() - this->blockSize);
-        idToLabelMapping.shrink_to_fit();
-        resizeLabelLookup(idToLabelMapping.size());
+        resizeIndexCommon(indexCapacity() + this->blockSize);
+    }
+
+    void shrinkByBlock() {
+        assert(indexCapacity() >= this->blockSize);
+        assert(indexCapacity() % this->blockSize == 0);
+        assert(dynamic_cast<DataBlocksContainer *>(this->vectors)->numBlocks() ==
+               indexSize() / this->blockSize);
+
+        if (indexSize() == 0) {
+            resizeIndexCommon(0);
+        } else if (indexCapacity() >= (indexSize() + 2 * this->blockSize)) {
+
+            assert(indexCapacity() == idToLabelMapping.capacity());
+            assert(idToLabelMapping.size() == idToLabelMapping.capacity());
+            assert(dynamic_cast<DataBlocksContainer *>(this->vectors)->size() +
+                       2 * this->blockSize ==
+                   idToLabelMapping.capacity());
+            resizeIndexCommon(indexCapacity() - this->blockSize);
+        }
     }
 
     void setVectorLabel(idType id, labelType new_label) { idToLabelMapping.at(id) = new_label; }
@@ -143,14 +168,15 @@ BruteForceIndex<DataType, DistType>::BruteForceIndex(
 
 template <typename DataType, typename DistType>
 void BruteForceIndex<DataType, DistType>::appendVector(const void *vector_data, labelType label) {
+    // Resize the index meta data structures if needed
+    if (indexSize() >= indexCapacity()) {
+        growByBlock();
+    }
+
     auto processed_blob = this->preprocessForStorage(vector_data);
     // Give the vector new id and increase count.
     idType id = this->count++;
 
-    // Resize the index meta data structures if needed
-    if (indexSize() > indexCapacity()) {
-        growByBlock();
-    }
     // add vector data to vector raw data container
     this->vectors->addElement(processed_blob.get(), id);
 
@@ -199,7 +225,7 @@ size_t BruteForceIndex<DataType, DistType>::indexSize() const {
 
 template <typename DataType, typename DistType>
 size_t BruteForceIndex<DataType, DistType>::indexCapacity() const {
-    return this->idToLabelMapping.size();
+    return this->idToLabelMapping.capacity();
 }
 
 template <typename DataType, typename DistType>

src/VecSim/algorithms/hnsw/hnsw.h

@@ -235,6 +235,11 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,
     double getEpsilon() const;
     size_t indexSize() const override;
     size_t indexCapacity() const override;
+    /**
+     * Checks if the index capacity is full to hint the caller a resize is needed.
+     * @note Must be called with indexDataGuard locked.
+     */
+    size_t isCapacityFull() const;
     size_t getEfConstruction() const;
     size_t getM() const;
     size_t getMaxLevel() const;
@@ -349,6 +354,11 @@ size_t HNSWIndex<DataType, DistType>::indexCapacity() const {
     return this->maxElements;
 }
 
+template <typename DataType, typename DistType>
+size_t HNSWIndex<DataType, DistType>::isCapacityFull() const {
+    return indexSize() == this->maxElements;
+}
+
 template <typename DataType, typename DistType>
 size_t HNSWIndex<DataType, DistType>::getEfConstruction() const {
     return this->efConstruction;
@@ -1281,31 +1291,59 @@ template <typename DataType, typename DistType>
 void HNSWIndex<DataType, DistType>::resizeIndexCommon(size_t new_max_elements) {
     assert(new_max_elements % this->blockSize == 0 &&
            "new_max_elements must be a multiple of blockSize");
-    this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
-              "Updating HNSW index capacity from %zu to %zu", this->maxElements, new_max_elements);
+    this->log(VecSimCommonStrings::LOG_VERBOSE_STRING, "Resizing HNSW index from %zu to %zu",
+              idToMetaData.capacity(), new_max_elements);
     resizeLabelLookup(new_max_elements);
     visitedNodesHandlerPool.resize(new_max_elements);
+    assert(idToMetaData.capacity() == idToMetaData.size());
     idToMetaData.resize(new_max_elements);
     idToMetaData.shrink_to_fit();
-
-    maxElements = new_max_elements;
+    assert(idToMetaData.capacity() == idToMetaData.size());
 }
 
 template <typename DataType, typename DistType>
 void HNSWIndex<DataType, DistType>::growByBlock() {
-    size_t new_max_elements = maxElements + this->blockSize;
+    assert(this->maxElements % this->blockSize == 0);
+    assert(this->maxElements == indexSize());
+    assert(graphDataBlocks.size() == this->maxElements / this->blockSize);
+    assert(idToMetaData.capacity() == maxElements ||
+           idToMetaData.capacity() == maxElements + this->blockSize);
+
+    this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
+              "Updating HNSW index capacity from %zu to %zu", maxElements,
+              maxElements + this->blockSize);
+    maxElements += this->blockSize;
+
     graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize, this->allocator);
 
-    resizeIndexCommon(new_max_elements);
+    if (idToMetaData.capacity() == indexSize()) {
+        resizeIndexCommon(maxElements);
+    }
 }
 
 template <typename DataType, typename DistType>
 void HNSWIndex<DataType, DistType>::shrinkByBlock() {
-    assert(maxElements >= this->blockSize);
-    size_t new_max_elements = maxElements - this->blockSize;
-    graphDataBlocks.pop_back();
+    assert(this->maxElements >= this->blockSize);
+    assert(this->maxElements % this->blockSize == 0);
+
+    if (indexSize() % this->blockSize == 0) {
+        this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
+                  "Updating HNSW index capacity from %zu to %zu", maxElements,
+                  maxElements - this->blockSize);
+        graphDataBlocks.pop_back();
+        assert(graphDataBlocks.size() == indexSize() / this->blockSize);
+
+        // assuming idToMetaData reflects the capacity of the heavy reallocation containers.
+        if (indexSize() == 0) {
+            resizeIndexCommon(0);
+        } else if (idToMetaData.capacity() >= (indexSize() + 2 * this->blockSize)) {
+            assert(this->maxElements + this->blockSize == idToMetaData.capacity());
+            resizeIndexCommon(idToMetaData.capacity() - this->blockSize);
+        }
 
-    resizeIndexCommon(new_max_elements);
+        // Take the lower bound into account.
+        maxElements -= this->blockSize;
+    }
 }
 
 template <typename DataType, typename DistType>
@@ -1660,9 +1698,7 @@ void HNSWIndex<DataType, DistType>::removeAndSwap(idType internalId) {
     // If we need to free a complete block and there is at least one block between the
     // capacity and the size.
     this->vectors->removeElement(curElementCount);
-    if (curElementCount % this->blockSize == 0) {
-        shrinkByBlock();
-    }
+    shrinkByBlock();
 }
 
 template <typename DataType, typename DistType>
@@ -1738,6 +1774,9 @@ void HNSWIndex<DataType, DistType>::removeVectorInPlace(const idType element_int
 template <typename DataType, typename DistType>
 HNSWAddVectorState HNSWIndex<DataType, DistType>::storeNewElement(labelType label,
                                                                   const void *vector_data) {
+    if (isCapacityFull()) {
+        growByBlock();
+    }
     HNSWAddVectorState state{};
 
     // Choose randomly the maximum level in which the new element will be in the index.
@@ -1765,14 +1804,6 @@ HNSWAddVectorState HNSWIndex<DataType, DistType>::storeNewElement(labelType labe
         throw e;
     }
 
-    if (indexSize() > indexCapacity()) {
-        growByBlock();
-    } else if (state.newElementId % this->blockSize == 0) {
-        // If we had an initial capacity, we might have to allocate new blocks for the graph data.
-        this->graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize,
-                                           this->allocator);
-    }
-
     // Insert the new element to the data block
     this->vectors->addElement(vector_data, state.newElementId);
     this->graphDataBlocks.back().addElement(cur_egd);

src/VecSim/algorithms/hnsw/hnsw_tiered.h

@@ -314,7 +314,7 @@ template <typename DataType, typename DistType>
 void TieredHNSWIndex<DataType, DistType>::executeReadySwapJobs(size_t maxJobsToRun) {
 
     // Execute swap jobs - acquire hnsw write lock.
-    this->mainIndexGuard.lock();
+    this->lockMainIndexGuard();
     TIERED_LOG(VecSimCommonStrings::LOG_VERBOSE_STRING,
                "Tiered HNSW index GC: there are %zu ready swap jobs. Start executing %zu swap jobs",
                readySwapJobs, std::min(readySwapJobs, maxJobsToRun));
@@ -339,7 +339,7 @@ void TieredHNSWIndex<DataType, DistType>::executeReadySwapJobs(size_t maxJobsToR
     readySwapJobs -= idsToRemove.size();
     TIERED_LOG(VecSimCommonStrings::LOG_VERBOSE_STRING,
                "Tiered HNSW index GC: done executing %zu swap jobs", idsToRemove.size());
-    this->mainIndexGuard.unlock();
+    this->unlockMainIndexGuard();
 }
 
 template <typename DataType, typename DistType>
@@ -437,11 +437,11 @@ void TieredHNSWIndex<DataType, DistType>::insertVectorToHNSW(
     this->mainIndexGuard.lock_shared();
     hnsw_index->lockIndexDataGuard();
     // Check if resizing is needed for HNSW index (requires write lock).
-    if (hnsw_index->indexCapacity() == hnsw_index->indexSize()) {
+    if (hnsw_index->isCapacityFull()) {
         // Release the inner HNSW data lock before we re-acquire the global HNSW lock.
         this->mainIndexGuard.unlock_shared();
         hnsw_index->unlockIndexDataGuard();
-        this->mainIndexGuard.lock();
+        this->lockMainIndexGuard();
         hnsw_index->lockIndexDataGuard();
 
         // Hold the index data lock while we store the new element. If the new node's max level is
@@ -466,7 +466,7 @@ void TieredHNSWIndex<DataType, DistType>::insertVectorToHNSW(
         if (state.elementMaxLevel > state.currMaxLevel) {
             hnsw_index->unlockIndexDataGuard();
         }
-        this->mainIndexGuard.unlock();
+        this->unlockMainIndexGuard();
     } else {
         // Do the same as above except for changing the capacity, but with *shared* lock held:
         // Hold the index data lock while we store the new element. If the new node's max level is
@@ -713,9 +713,9 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
         auto storage_blob = this->frontendIndex->preprocessForStorage(blob);
         // Insert the vector to the HNSW index. Internally, we will never have to overwrite the
         // label since we already checked it outside.
-        this->mainIndexGuard.lock();
+        this->lockMainIndexGuard();
         hnsw_index->addVector(storage_blob.get(), label);
-        this->mainIndexGuard.unlock();
+        this->unlockMainIndexGuard();
         return ret;
     }
     if (this->frontendIndex->indexSize() >= this->flatBufferLimit) {
@@ -841,9 +841,9 @@ int TieredHNSWIndex<DataType, DistType>::deleteVector(labelType label) {
         }
     } else {
         // delete in place.
-        this->mainIndexGuard.lock();
+        this->lockMainIndexGuard();
         num_deleted_vectors += this->deleteLabelFromHNSWInplace(label);
-        this->mainIndexGuard.unlock();
+        this->unlockMainIndexGuard();
     }
 
     return num_deleted_vectors;

src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h

@@ -63,6 +63,7 @@ INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteBothAsyncAndInplaceMulti_
 INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteInplaceMultiSwapId_Test)
 INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteInplaceAvoidUpdatedMarkedDeleted_Test)
 INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_switchDeleteModes_Test)
+INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_HNSWResize_Test)
 
 friend class CommonAPITest_SearchDifferentScores_Test;
 friend class BF16TieredTest;

src/VecSim/containers/data_blocks_container.cpp

@@ -64,6 +64,8 @@ std::unique_ptr<RawDataContainer::Iterator> DataBlocksContainer::getIterator() c
     return std::make_unique<DataBlocksContainer::Iterator>(*this);
 }
 
+size_t DataBlocksContainer::numBlocks() const { return this->blocks.size(); }
+
 #ifdef BUILD_TESTS
 void DataBlocksContainer::saveVectorsData(std::ostream &output) const {
     // Save data blocks
@@ -114,8 +116,6 @@ void DataBlocksContainer::restoreBlocks(std::istream &input, size_t num_vectors,
 
 void DataBlocksContainer::shrinkToFit() { this->blocks.shrink_to_fit(); }
 
-size_t DataBlocksContainer::numBlocks() const { return this->blocks.size(); }
-
 #endif
 /********************************** Iterator API ************************************************/
 

src/VecSim/containers/data_blocks_container.h

@@ -28,8 +28,10 @@ class DataBlocksContainer : public VecsimBaseObject, public RawDataContainer {
                         std::shared_ptr<VecSimAllocator> allocator, unsigned char alignment = 0);
     ~DataBlocksContainer();
 
+    // Number of elements in the container.
     size_t size() const override;
 
+    // Number of blocks allocated.
     size_t capacity() const;
 
     size_t blockSize() const;
@@ -46,13 +48,13 @@ class DataBlocksContainer : public VecsimBaseObject, public RawDataContainer {
 
     std::unique_ptr<RawDataContainer::Iterator> getIterator() const override;
 
+    size_t numBlocks() const;
 #ifdef BUILD_TESTS
     void saveVectorsData(std::ostream &output) const override;
     // Use that in deserialization when file was created with old version (v3) that serialized
     // the blocks themselves and not just thw raw vector data.
     void restoreBlocks(std::istream &input, size_t num_vectors, Serializer::EncodingVersion);
     void shrinkToFit();
-    size_t numBlocks() const;
 #endif
 
     class Iterator : public RawDataContainer::Iterator {

src/VecSim/vec_sim_tiered_index.h

@@ -50,7 +50,17 @@ class VecSimTieredIndex : public VecSimIndexInterface {
 
     mutable std::shared_mutex flatIndexGuard;
     mutable std::shared_mutex mainIndexGuard;
+    void lockMainIndexGuard() const {
+        mainIndexGuard.lock();
+#ifdef BUILD_TESTS
+        mainIndexGuard_write_lock_count++;
+#endif
+    }
 
+    void unlockMainIndexGuard() const { mainIndexGuard.unlock(); }
+#ifdef BUILD_TESTS
+    mutable std::atomic_int mainIndexGuard_write_lock_count = 0;
+#endif
     size_t flatBufferLimit;
 
     void submitSingleJob(AsyncJob *job) {
@@ -89,6 +99,7 @@ class VecSimTieredIndex : public VecSimIndexInterface {
 
 #ifdef BUILD_TESTS
 public:
+    int getMainIndexGuardWriteLockCount() const { return mainIndexGuard_write_lock_count; }
 #endif
     // For both topK and range, Use withSet=false if you can guarantee that shared ids between the
     // two lists will also have identical scores. In this case, any duplicates will naturally align