Add index lock

Beihao-Zhou · Aug 16, 2024 · 34c09d7 · 34c09d7
1 parent a86d317
commit 34c09d7
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/src/search/index_info.h b/src/search/index_info.h
@@ -23,6 +23,7 @@
 #include <algorithm>
 #include <map>
 #include <memory>
+#include <mutex>
 #include <string>
 #include <utility>
 
@@ -52,21 +53,28 @@ struct FieldInfo {
 
 struct IndexInfo {
   using FieldMap = std::map<std::string, FieldInfo>;
+  using MutexMap = std::map<std::string, std::mutex>;
 
   std::string name;
   redis::IndexMetadata metadata;
   FieldMap fields;
+  mutable MutexMap field_mutexes;
   redis::IndexPrefixes prefixes;
   std::string ns;
 
   IndexInfo(std::string name, redis::IndexMetadata metadata, std::string ns)
       : name(std::move(name)), metadata(std::move(metadata)), ns(std::move(ns)) {}
 
   void Add(FieldInfo &&field) {
-    const auto &name = field.name;
+    auto name = field.name;
     field.index = this;
     fields.emplace(name, std::move(field));
+    field_mutexes.emplace(std::piecewise_construct, std::make_tuple(name), std::make_tuple());
   }
+
+  void LockField(const std::string &field_name) const { field_mutexes.at(field_name).lock(); }
+
+  void UnLockField(const std::string &field_name) const { field_mutexes.at(field_name).unlock(); }
 };
 
 struct IndexMap : std::map<std::string, std::unique_ptr<IndexInfo>> {

diff --git a/src/search/indexer.cc b/src/search/indexer.cc
@@ -304,19 +304,23 @@ Status IndexUpdater::UpdateIndex(const std::string &field, std::string_view key,
     return {Status::NotOK, "No such field to do index updating"};
   }
 
+  info->LockField(field);
+
   auto *metadata = iter->second.metadata.get();
   SearchKey search_key(info->ns, info->name, field);
+  Status status;
   if (auto tag = dynamic_cast<TagFieldMetadata *>(metadata)) {
-    GET_OR_RET(UpdateTagIndex(key, original, current, search_key, tag));
+    status = UpdateTagIndex(key, original, current, search_key, tag);
   } else if (auto numeric [[maybe_unused]] = dynamic_cast<NumericFieldMetadata *>(metadata)) {
-    GET_OR_RET(UpdateNumericIndex(key, original, current, search_key, numeric));
+    status = UpdateNumericIndex(key, original, current, search_key, numeric);
   } else if (auto vector = dynamic_cast<HnswVectorFieldMetadata *>(metadata)) {
-    GET_OR_RET(UpdateHnswVectorIndex(key, original, current, search_key, vector));
+    status = UpdateHnswVectorIndex(key, original, current, search_key, vector);
   } else {
-    return {Status::NotOK, "Unexpected field type"};
+    status = {Status::NotOK, "Unexpected field type"};
   }
 
-  return Status::OK();
+  info->UnLockField(field);
+  return status;
 }
 
 Status IndexUpdater::Update(const FieldValues &original, std::string_view key) const {