Merge branch 'main' into abhijat/feat/implement-chunk-tag-load

Author: abhijat

src/core/search/hnsw_index.cc

@@ -132,35 +132,13 @@ struct HnswlibAdapter {
   HnswIndexMetadata GetMetadata() const {
     MRMWMutexLock lock(&mrmw_mutex_, MRMWMutex::LockMode::kReadLock);
     HnswIndexMetadata metadata;
-    metadata.max_elements = world_.max_elements_;
-    metadata.cur_element_count = world_.cur_element_count.load();
-    metadata.maxlevel = world_.maxlevel_;
     metadata.enterpoint_node = world_.enterpoint_node_;
     return metadata;
   }
 
-  void SetMetadata(const HnswIndexMetadata& metadata) {
-    MRMWMutexLock lock(&mrmw_mutex_, MRMWMutex::LockMode::kWriteLock);
-    absl::WriterMutexLock resize_lock(&resize_mutex_);
-
-    // SetMetadata is only called during deserialization before the index is used.
-    // Assert the index is empty to ensure no concurrent operations are possible.
-    DCHECK_EQ(world_.cur_element_count.load(), 0u)
-        << "SetMetadata should only be called on an empty index during deserialization";
-
-    // Runtime check for release builds to prevent silent corruption
-    if (world_.cur_element_count.load() != 0) {
-      LOG(ERROR) << "SetMetadata called on non-empty HNSW index with "
-                 << world_.cur_element_count.load() << " elements, ignoring";
-      return;
-    }
-
-    // Pre-allocate capacity based on expected element count, but don't set cur_element_count.
-    // cur_element_count will be set by RestoreFromNodes when the actual nodes are restored.
-    if (world_.max_elements_ < metadata.cur_element_count) {
-      world_.resizeIndex(metadata.cur_element_count);
-    }
-    // Note: Don't set cur_element_count here - RestoreFromNodes will set it after restoring nodes.
+  int GetMaxLevel() const {
+    MRMWMutexLock lock(&mrmw_mutex_, MRMWMutex::LockMode::kReadLock);
+    return world_.maxlevel_;
   }
 
   size_t GetNodeCount() const {
@@ -280,31 +258,41 @@ struct HnswlibAdapter {
   }
 
  public:
-  // Restore HNSW graph structure from serialized nodes with metadata
-  void RestoreFromNodes(const std::vector<HnswNodeData>& nodes, const HnswIndexMetadata& metadata) {
+  // Restore HNSW graph structure from serialized nodes with metadata.
+  // Returns false if the input is inconsistent (e.g. entry point not in node set) —
+  // caller should fall back to rebuilding the index from the keyspace.
+  bool RestoreFromNodes(const std::vector<HnswNodeData>& nodes, const HnswIndexMetadata& metadata) {
     MRMWMutexLock lock(&mrmw_mutex_, MRMWMutex::LockMode::kWriteLock);
     absl::WriterMutexLock resize_lock(&resize_mutex_);
 
     if (nodes.empty()) {
-      return;
+      return true;
     }
 
     // RestoreFromNodes is only called during deserialization on a freshly created index.
     // Assert the index is empty to prevent memory leaks from double-allocation of linkLists_.
     DCHECK_EQ(world_.cur_element_count.load(), 0u)
         << "RestoreFromNodes should only be called on an empty index during deserialization";
 
-    // Ensure we have enough capacity.
-    // Metadata may have been captured before the snapshot read-lock, so
-    // cur_element_count can be smaller than actual node internal_ids when
-    // concurrent writes happen.  Compute the real requirement from nodes.
+    // hnswlib pairs enterpoint_node_ with maxlevel_; node levels are immutable after
+    // creation, so the entry point's level in the serialized set equals the live
+    // maxlevel at metadata capture. max(node.level) would risk OOB reads when a
+    // concurrent Add raised maxlevel between capture and node serialization.
     size_t max_internal_id = 0;
+    int entrypoint_level = -1;
     for (const auto& node : nodes) {
       max_internal_id = std::max<size_t>(max_internal_id, node.internal_id);
+      if (node.internal_id == metadata.enterpoint_node)
+        entrypoint_level = node.level;
     }
-    size_t required_capacity = std::max(metadata.cur_element_count, max_internal_id + 1);
-    if (world_.max_elements_ < required_capacity) {
-      world_.resizeIndex(required_capacity);
+    if (entrypoint_level < 0) {
+      LOG(ERROR) << "HNSW restore: entry point internal_id=" << metadata.enterpoint_node
+                 << " not present in serialized node set (" << nodes.size()
+                 << " nodes); skipping restore — index will be rebuilt from the keyspace";
+      return false;
+    }
+    if (world_.max_elements_ < max_internal_id + 1) {
+      world_.resizeIndex(max_internal_id + 1);
     }
 
     // Restore each node - directly set up memory and fields
@@ -378,12 +366,13 @@ struct HnswlibAdapter {
     }
 
     // Set the metadata for the graph
-    world_.maxlevel_ = metadata.maxlevel;
+    world_.maxlevel_ = entrypoint_level;
     world_.enterpoint_node_ = metadata.enterpoint_node;
 
     VLOG(1) << "Restored HNSW index with " << restored_count
-            << " nodes, maxlevel=" << metadata.maxlevel
+            << " nodes, maxlevel=" << entrypoint_level
             << ", enterpoint=" << metadata.enterpoint_node;
+    return true;
   }
 
   // Update vector data for an existing node (used after RestoreFromNodes).
@@ -502,8 +491,8 @@ HnswIndexMetadata HnswVectorIndex::GetMetadata() const {
   return adapter_->GetMetadata();
 }
 
-void HnswVectorIndex::SetMetadata(const HnswIndexMetadata& metadata) {
-  adapter_->SetMetadata(metadata);
+int HnswVectorIndex::GetMaxLevel() const {
+  return adapter_->GetMaxLevel();
 }
 
 size_t HnswVectorIndex::GetNodeCount() const {
@@ -514,9 +503,9 @@ std::vector<HnswNodeData> HnswVectorIndex::GetNodesRange(size_t start, size_t en
   return adapter_->GetNodesRange(start, end);
 }
 
-void HnswVectorIndex::RestoreFromNodes(const std::vector<HnswNodeData>& nodes,
+bool HnswVectorIndex::RestoreFromNodes(const std::vector<HnswNodeData>& nodes,
                                        const HnswIndexMetadata& metadata) {
-  adapter_->RestoreFromNodes(nodes, metadata);
+  return adapter_->RestoreFromNodes(nodes, metadata);
 }
 
 bool HnswVectorIndex::UpdateVectorData(GlobalDocId id, const DocumentAccessor& doc,

src/core/search/hnsw_index.h

@@ -11,16 +11,12 @@
 
 namespace dfly::search {
 
-// Metadata structure for HNSW index serialization
-// Contains the key parameters needed to restore the index state
+// Wire format for HNSW index AUX. Only the entry point is persisted: capacity is
+// derived from max(internal_id)+1 in the node set and maxlevel from the entry-point
+// node's level (hnswlib pairs enterpoint_node_ with maxlevel_, and node levels are
+// immutable after creation).
 struct HnswIndexMetadata {
-  size_t max_elements = 0;  // Maximum number of elements the index can hold
-  // Note: cur_element_count may be smaller than actual node count during concurrent writes,
-  // so we compute the real requirement from nodes during restoration.
-  // TODO: consider removing it from metadata and rely entirely on node data for restoration.
-  size_t cur_element_count = 0;  // Current number of elements in the index
-  int maxlevel = -1;             // Maximum level of the graph
-  size_t enterpoint_node = 0;    // Entry point node for the graph
+  size_t enterpoint_node = 0;
 };
 
 // Node data structure for HNSW serialization
@@ -75,8 +71,9 @@ class HnswVectorIndex {
   // Get metadata for serialization
   HnswIndexMetadata GetMetadata() const;
 
-  // Set metadata (used during restoration)
-  void SetMetadata(const HnswIndexMetadata& metadata);
+  // Current graph maxlevel_. Exposed for introspection and tests that need to
+  // verify invariants preserved by RestoreFromNodes (entry point must sit at maxlevel).
+  int GetMaxLevel() const;
 
   // Get total number of nodes in the index
   size_t GetNodeCount() const;
@@ -85,10 +82,12 @@ class HnswVectorIndex {
   // Returns vector of node data for serialization
   std::vector<HnswNodeData> GetNodesRange(size_t start, size_t end) const;
 
-  // Restore graph structure from serialized nodes with metadata
-  // This restores the HNSW graph links but NOT the vector data
-  // Vector data must be populated separately via UpdateVectorData
-  void RestoreFromNodes(const std::vector<HnswNodeData>& nodes, const HnswIndexMetadata& metadata);
+  // Restore graph structure from serialized nodes with metadata.
+  // Restores links only; vector data must be populated separately via UpdateVectorData.
+  // Returns false if the metadata is inconsistent with the node set (e.g. the entry
+  // point is missing from the serialized nodes) — caller should then leave the index
+  // empty and let the higher-level rebuild path repopulate it from the keyspace.
+  bool RestoreFromNodes(const std::vector<HnswNodeData>& nodes, const HnswIndexMetadata& metadata);
 
   // Update vector data for an existing node (used after RestoreFromNodes)
   // This populates the vector data for a node that already has graph links

src/core/search/search_test.cc

@@ -1072,6 +1072,30 @@ TEST_F(KnnTest, AutoResize) {
   EXPECT_EQ(indices.GetAllDocs().size(), 100);
 }
 
+// Seeds the given HNSW index with `n` deterministic random vectors of dim `dim` using
+// the given RNG seed. Returns the owning MockedDocuments so the caller can pass them
+// back to UpdateVectorData after a restore. Used by the serialization/restore tests.
+inline vector<MockedDocument> SeedHnswIndex(HnswVectorIndex& index, size_t n, size_t dim,
+                                            uint32_t rng_seed) {
+  vector<MockedDocument> docs(n);
+  std::mt19937 rng(rng_seed);
+  std::uniform_real_distribution<float> dist(0.0f, 1.0f);
+  for (size_t i = 0; i < n; i++) {
+    vector<float> coords(dim);
+    for (size_t d = 0; d < dim; d++)
+      coords[d] = dist(rng);
+    docs[i] = MockedDocument::Map{{"vec", ToBytes(absl::MakeConstSpan(coords))}};
+    index.Add(i, docs[i], "vec");
+  }
+  return docs;
+}
+
+// Snapshots all nodes from the index under its read lock.
+inline vector<HnswNodeData> SnapshotHnswNodes(const HnswVectorIndex& index) {
+  auto lock = index.GetReadLock();
+  return index.GetNodesRange(0, index.GetNodeCount());
+}
+
 // Parameterized HNSW serialization round-trip test.
 // Parameters: {num_elements, dim, similarity}
 struct HnswSerParam {
@@ -1108,27 +1132,12 @@ TEST_P(HnswSerializationTest, RoundTrip) {
   params.hnsw_ef_construction = 200;
 
   HnswVectorIndex original(params, /*copy_vector=*/true);
+  vector<MockedDocument> docs = SeedHnswIndex(original, num_elements, dim, /*rng_seed=*/42);
 
-  std::mt19937 rng(42);
-  std::uniform_real_distribution<float> dist(0.0f, 1.0f);
-  vector<MockedDocument> docs(num_elements);
-  for (size_t i = 0; i < num_elements; i++) {
-    vector<float> coords(dim);
-    for (size_t d = 0; d < dim; d++)
-      coords[d] = dist(rng);
-    docs[i] = MockedDocument::Map{{"vec", ToBytes(absl::MakeConstSpan(coords))}};
-    original.Add(i, docs[i], "vec");
-  }
-
-  // Serialize
   auto metadata = original.GetMetadata();
-  ASSERT_EQ(metadata.cur_element_count, num_elements);
+  ASSERT_EQ(original.GetNodeCount(), num_elements);
 
-  std::vector<HnswNodeData> nodes;
-  {
-    auto lock = original.GetReadLock();
-    nodes = original.GetNodesRange(0, metadata.cur_element_count);
-  }
+  std::vector<HnswNodeData> nodes = SnapshotHnswNodes(original);
   ASSERT_EQ(nodes.size(), num_elements);
 
   // Verify node data integrity
@@ -1139,8 +1148,7 @@ TEST_P(HnswSerializationTest, RoundTrip) {
 
   // Deserialize into a fresh index
   HnswVectorIndex restored(params, /*copy_vector=*/true);
-  restored.SetMetadata(metadata);
-  restored.RestoreFromNodes(nodes, metadata);
+  ASSERT_TRUE(restored.RestoreFromNodes(nodes, metadata));
 
   // Before UpdateVectorData, all nodes must be marked deleted.
   // KNN should safely return empty results (no crash from nullptr dereference).
@@ -1153,17 +1161,16 @@ TEST_P(HnswSerializationTest, RoundTrip) {
   for (size_t i = 0; i < num_elements; i++)
     restored.UpdateVectorData(i, docs[i], "vec");
 
-  // Metadata must match
   auto rm = restored.GetMetadata();
-  EXPECT_EQ(rm.cur_element_count, metadata.cur_element_count);
-  EXPECT_EQ(rm.maxlevel, metadata.maxlevel);
+  EXPECT_EQ(restored.GetNodeCount(), num_elements);
   EXPECT_EQ(rm.enterpoint_node, metadata.enterpoint_node);
+  EXPECT_EQ(restored.GetMaxLevel(), original.GetMaxLevel());
 
   // Graph links must be identical
   std::vector<HnswNodeData> restored_nodes;
   {
     auto lock = restored.GetReadLock();
-    restored_nodes = restored.GetNodesRange(0, rm.cur_element_count);
+    restored_nodes = restored.GetNodesRange(0, restored.GetNodeCount());
   }
   ASSERT_EQ(restored_nodes.size(), nodes.size());
   for (size_t i = 0; i < nodes.size(); i++) {
@@ -1209,6 +1216,76 @@ TEST_P(HnswSerializationTest, RoundTrip) {
   }
 }
 
+// Regression for the save-side race where an Add raises maxlevel between metadata
+// capture and node serialization (see RestoreFromNodes for the rationale). Simulated
+// by forging metadata with a low-level entry point against a multi-level node set;
+// expects maxlevel_ to clamp to the entry point's level rather than max(node.level).
+TEST(HnswRestoreInvariant, MaxLevelClampedToEntryPointLevel) {
+  constexpr size_t kDim = 8;
+  constexpr size_t kN = 100;
+
+  InitTLSearchMR(PMR_NS::get_default_resource());
+  absl::Cleanup cleanup = [] { InitTLSearchMR(nullptr); };
+
+  SchemaField::VectorParams params;
+  params.use_hnsw = true;
+  params.dim = kDim;
+  params.sim = VectorSimilarity::L2;
+  params.capacity = kN;
+  params.hnsw_m = 16;
+  params.hnsw_ef_construction = 200;
+
+  HnswVectorIndex original(params, /*copy_vector=*/true);
+  SeedHnswIndex(original, kN, kDim, /*rng_seed=*/42);
+  std::vector<HnswNodeData> nodes = SnapshotHnswNodes(original);
+
+  int global_max_level = -1;
+  std::optional<uint32_t> low_level_internal_id;
+  for (const auto& n : nodes) {
+    global_max_level = std::max(global_max_level, n.level);
+    if (!low_level_internal_id && n.level == 0)
+      low_level_internal_id = n.internal_id;
+  }
+  ASSERT_GT(global_max_level, 0) << "test setup: need a multi-level graph";
+  ASSERT_TRUE(low_level_internal_id.has_value()) << "test setup: need a level-0 node";
+
+  HnswIndexMetadata forged_metadata{.enterpoint_node = *low_level_internal_id};
+
+  HnswVectorIndex restored(params, /*copy_vector=*/true);
+  ASSERT_TRUE(restored.RestoreFromNodes(nodes, forged_metadata));
+
+  EXPECT_EQ(restored.GetMaxLevel(), 0)
+      << "maxlevel_ must equal entry-point level; got " << restored.GetMaxLevel()
+      << " while node set max level=" << global_max_level;
+}
+
+// Malformed/mismatched metadata (entry point not in serialized node set) must
+// fail restoration gracefully — returning false — instead of SIGABRT'ing via
+// CHECK. Callers then rebuild the index from the keyspace.
+TEST(HnswRestoreInvariant, MissingEntrypointFailsGracefully) {
+  constexpr size_t kDim = 4;
+  constexpr size_t kN = 10;
+
+  InitTLSearchMR(PMR_NS::get_default_resource());
+  absl::Cleanup cleanup = [] { InitTLSearchMR(nullptr); };
+
+  SchemaField::VectorParams params;
+  params.use_hnsw = true;
+  params.dim = kDim;
+  params.sim = VectorSimilarity::L2;
+  params.capacity = kN;
+  params.hnsw_m = 16;
+  params.hnsw_ef_construction = 200;
+
+  HnswVectorIndex original(params, /*copy_vector=*/true);
+  SeedHnswIndex(original, kN, kDim, /*rng_seed=*/7);
+  std::vector<HnswNodeData> nodes = SnapshotHnswNodes(original);
+
+  HnswIndexMetadata bad_metadata{.enterpoint_node = 999999};  // well past any real id
+  HnswVectorIndex restored(params, /*copy_vector=*/true);
+  EXPECT_FALSE(restored.RestoreFromNodes(nodes, bad_metadata));
+}
+
 // Regression: in borrowed mode (copy_vector=false), Remove marks the node deleted
 // but hnswlib still traverses it and dereferences its data pointer.  If the external
 // data is freed (as happens after DEL), the pointer dangles.  The fix in DoRemove