Reclaim memory when doing changes to the graph in the hnsw index.

5 files changed, 104 insertions, 8 deletions

diff --git a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
index 761d403ac6c..1f3ad7c2fec 100644
--- a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
+++ b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
@@ -134,6 +134,9 @@ public:
     void trim_hold_lists(generation_t first_used_gen) override {
         _trim_gen = first_used_gen;
     }
+    vespalib::MemoryUsage memory_usage() const override {
+        return vespalib::MemoryUsage();
+    }
     std::vector<Neighbor> find_top_k(uint32_t k, vespalib::tensor::TypedCells vector, uint32_t explore_k) const override {
         (void) k;
         (void) vector;
diff --git a/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp b/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
index 52f45860c1e..37c4d02017f 100644
--- a/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
+++ b/searchlib/src/tests/tensor/hnsw_index/hnsw_index_test.cpp
@@ -6,11 +6,14 @@
 #include <vespa/searchlib/tensor/hnsw_index.h>
 #include <vespa/searchlib/tensor/random_level_generator.h>
 #include <vespa/vespalib/gtest/gtest.h>
+#include <vespa/vespalib/util/generationhandler.h>
 #include <vector>
 
 #include <vespa/log/log.h>
 LOG_SETUP("hnsw_index_test");
 
+using vespalib::GenerationHandler;
+using vespalib::MemoryUsage;
 using namespace search::tensor;
 
 template <typename FloatType>
@@ -49,11 +52,13 @@ class HnswIndexTest : public ::testing::Test {
 public:
     FloatVectors vectors;
     LevelGenerator* level_generator;
+    GenerationHandler gen_handler;
     HnswIndexUP index;
 
     HnswIndexTest()
         : vectors(),
           level_generator(),
+          gen_handler(),
           index()
     {
         vectors.set(1, {2, 2}).set(2, {3, 2}).set(3, {2, 3})
@@ -70,9 +75,23 @@ public:
     void add_document(uint32_t docid, uint32_t max_level = 0) {
         level_generator->level = max_level;
         index->add_document(docid);
+        commit();
     }
     void remove_document(uint32_t docid) {
         index->remove_document(docid);
+        commit();
+    }
+    void commit() {
+        index->transfer_hold_lists(gen_handler.getCurrentGeneration());
+        gen_handler.incGeneration();
+        gen_handler.updateFirstUsedGeneration();
+        index->trim_hold_lists(gen_handler.getFirstUsedGeneration());
+    }
+    GenerationHandler::Guard take_read_guard() {
+        return gen_handler.takeGuard();
+    }
+    MemoryUsage memory_usage() const {
+        return index->memory_usage();
     }
     void expect_entry_point(uint32_t exp_docid, uint32_t exp_level) {
         EXPECT_EQ(exp_docid, index->get_entry_docid());
@@ -83,6 +102,10 @@ public:
         ASSERT_EQ(1, node.size());
         EXPECT_EQ(exp_links, node.level(0));
     }
+    void expect_empty_level_0(uint32_t docid) {
+        auto node = index->get_node(docid);
+        EXPECT_TRUE(node.empty());
+    }
     void expect_levels(uint32_t docid, const HnswNode::LevelArray& exp_levels) {
         auto act_node = index->get_node(docid);
         ASSERT_EQ(exp_levels.size(), act_node.size());
@@ -299,5 +322,50 @@ TEST_F(HnswIndexTest, manual_insert)
     expect_levels(5, {{1,2}, {4}});
 }
 
+TEST_F(HnswIndexTest, memory_is_reclaimed_when_doing_changes_to_graph)
+{
+    init(true);
+
+    add_document(1);
+    add_document(3);
+    auto mem_1 = memory_usage();
+
+    add_document(2);
+    expect_level_0(1, {2,3});
+    expect_level_0(2, {1,3});
+    expect_level_0(3, {1,2});
+    auto mem_2 = memory_usage();
+    // We should use more memory with larger link arrays and extra document.
+    EXPECT_GT((mem_2.usedBytes() - mem_2.deadBytes()), (mem_1.usedBytes() - mem_1.deadBytes()));
+    EXPECT_EQ(0, mem_2.allocatedBytesOnHold());
+
+    remove_document(2);
+    expect_level_0(1, {3});
+    expect_empty_level_0(2);
+    expect_level_0(3, {1});
+    auto mem_3 = memory_usage();
+    // We end up in the same state as before document 2 was added and effectively use the same amount of memory.
+    EXPECT_EQ((mem_1.usedBytes() - mem_1.deadBytes()), (mem_3.usedBytes() - mem_3.deadBytes()));
+    EXPECT_EQ(0, mem_3.allocatedBytesOnHold());
+}
+
+TEST_F(HnswIndexTest, memory_is_put_on_hold_while_read_guard_is_held)
+{
+    init(true);
+
+    add_document(1);
+    add_document(3);
+    {
+        auto guard = take_read_guard();
+        add_document(2);
+        auto mem = memory_usage();
+        // As read guard is held memory to reclaim is put on hold
+        EXPECT_GT(mem.allocatedBytesOnHold(), 0);
+    }
+    commit();
+    auto mem = memory_usage();
+    EXPECT_EQ(0, mem.allocatedBytesOnHold());
+}
+
 GTEST_MAIN_RUN_ALL_TESTS()
 
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
index e0b999bf5c5..467e41d83fc 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
@@ -9,6 +9,8 @@
 
 namespace search::tensor {
 
+using search::datastore::EntryRef;
+
 namespace {
 
 // TODO: Move this to MemoryAllocator, with name PAGE_SIZE.
@@ -44,6 +46,9 @@ HnswIndex::max_links_for_level(uint32_t level) const
 uint32_t
 HnswIndex::make_node_for_document(uint32_t docid)
 {
+    // A document cannot be added twice.
+    assert(!_node_refs[docid].load_acquire().valid());
+
     uint32_t max_level = _level_generator->max_level();
     // TODO: Add capping on num_levels
     uint32_t num_levels = max_level + 1;
@@ -54,6 +59,15 @@ HnswIndex::make_node_for_document(uint32_t docid)
     return max_level;
 }
 
+void
+HnswIndex::remove_node_for_document(uint32_t docid)
+{
+    auto node_ref = _node_refs[docid].load_acquire();
+    _nodes.remove(node_ref);
+    EntryRef invalid;
+    _node_refs[docid].store_release(invalid);
+}
+
 HnswIndex::LevelArrayRef
 HnswIndex::get_level_array(uint32_t docid) const
 {
@@ -72,10 +86,12 @@ HnswIndex::get_link_array(uint32_t docid, uint32_t level) const
 void
 HnswIndex::set_link_array(uint32_t docid, uint32_t level, const LinkArrayRef& links)
 {
-    auto links_ref = _links.add(links);
+    auto new_links_ref = _links.add(links);
     auto node_ref = _node_refs[docid].load_acquire();
     auto levels = _nodes.get_writable(node_ref);
-    levels[level].store_release(links_ref);
+    auto old_links_ref = levels[level].load_acquire();
+    levels[level].store_release(new_links_ref);
+    _links.remove(old_links_ref);
 }
 
 bool
@@ -248,8 +264,6 @@ HnswIndex::add_document(uint32_t docid)
 {
     auto input = get_vector(docid);
     _node_refs.ensure_size(docid + 1, AtomicEntryRef());
-    // A document cannot be added twice.
-    assert(!_node_refs[docid].load_acquire().valid());
     int level = make_node_for_document(docid);
     if (_entry_docid == 0) {
         _entry_docid = docid;
@@ -306,8 +320,7 @@ HnswIndex::remove_document(uint32_t docid)
         _entry_docid = 0;
         _entry_level = -1;
     }
-    search::datastore::EntryRef invalid;
-    _node_refs[docid].store_release(invalid);
+    remove_node_for_document(docid);
 }
 
 void
@@ -328,6 +341,16 @@ HnswIndex::trim_hold_lists(generation_t first_used_gen)
     _links.trimHoldLists(first_used_gen);
 }
 
+vespalib::MemoryUsage
+HnswIndex::memory_usage() const
+{
+    vespalib::MemoryUsage result;
+    result.merge(_node_refs.getMemoryUsage());
+    result.merge(_nodes.getMemoryUsage());
+    result.merge(_links.getMemoryUsage());
+    return result;
+}
+
 struct NeighborsByDocId {
     bool operator() (const NearestNeighborIndex::Neighbor &lhs,
                      const NearestNeighborIndex::Neighbor &rhs)
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
index afbb3207fc2..89c45d6b50c 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
@@ -95,6 +95,7 @@ protected:
 
     uint32_t max_links_for_level(uint32_t level) const;
     uint32_t make_node_for_document(uint32_t docid);
+    void remove_node_for_document(uint32_t docid);
     LevelArrayRef get_level_array(uint32_t docid) const;
     LinkArrayRef get_link_array(uint32_t docid, uint32_t level) const;
     void set_link_array(uint32_t docid, uint32_t level, const LinkArrayRef& links);
@@ -138,12 +139,11 @@ public:
     void remove_document(uint32_t docid) override;
     void transfer_hold_lists(generation_t current_gen) override;
     void trim_hold_lists(generation_t first_used_gen) override;
+    vespalib::MemoryUsage memory_usage() const override;
     std::vector<Neighbor> find_top_k(uint32_t k, TypedCells vector, uint32_t explore_k) const override;
 
     FurthestPriQ top_k_candidates(const TypedCells &vector, uint32_t k) const;
 
-    // TODO: Add support for generation handling and cleanup (transfer_hold_lists, trim_hold_lists)
-
     uint32_t get_entry_docid() const { return _entry_docid; }
     uint32_t get_entry_level() const { return _entry_level; }
 
diff --git a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h
index 003daa3185e..bd98623bdd3 100644
--- a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h
+++ b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h
@@ -6,6 +6,7 @@
 #include <vector>
 #include <vespa/eval/tensor/dense/typed_cells.h>
 #include <vespa/vespalib/util/generationhandler.h>
+#include <vespa/vespalib/util/memoryusage.h>
 
 namespace search::tensor {
 
@@ -28,6 +29,7 @@ public:
     virtual void remove_document(uint32_t docid) = 0;
     virtual void transfer_hold_lists(generation_t current_gen) = 0;
     virtual void trim_hold_lists(generation_t first_used_gen) = 0;
+    virtual vespalib::MemoryUsage memory_usage() const = 0;
 
     virtual std::vector<Neighbor> find_top_k(uint32_t k,
                                              vespalib::tensor::TypedCells vector,