From 81a8173a38189431ee2469ce583c1753f1709cbe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E8=88=AA=E5=AE=87?= <3364451258@qq.com>
Date: Sun, 5 Jul 2026 01:34:36 +0800
Subject: [PATCH] fix: add thread-safe write lock to ChromaDB operations

---
 src/core/db.py     | 20 ++++++++++++++------
 src/core/ingest.py | 24 +++++++++++++-----------
 src/core/search.py | 15 ++++++++-------
 3 files changed, 35 insertions(+), 24 deletions(-)

diff --git a/src/core/db.py b/src/core/db.py
index 728f35d..51ccd25 100644
--- a/src/core/db.py
+++ b/src/core/db.py
@@ -1,24 +1,32 @@
 """ChromaDB 数据库层."""
+import threading
 import chromadb
 from chromadb.api.models.Collection import Collection
 
 
 class VectorDB:
-    """向量数据库封装."""
+    """线程安全的向量数据库封装."""
 
     def __init__(self, persist_dir: str = "./data"):
         self.client = chromadb.PersistentClient(path=persist_dir)
+        self._write_lock = threading.Lock()
+
+    @property
+    def write_lock(self) -> threading.Lock:
+        """获取写锁, 供外部在 add/delete/update 操作时使用."""
+        return self._write_lock
 
     def get_or_create_collection(self, name: str) -> Collection:
         """获取或创建 collection."""
         return self.client.get_or_create_collection(name=name)
 
     def delete_collection(self, name: str) -> None:
-        """删除 collection."""
-        try:
-            self.client.delete_collection(name=name)
-        except ValueError:
-            pass  # collection 不存在则忽略
+        """删除 collection (线程安全)."""
+        with self._write_lock:
+            try:
+                self.client.delete_collection(name=name)
+            except ValueError:
+                pass  # collection 不存在则忽略
 
     def close(self) -> None:
         """释放数据库连接."""
diff --git a/src/core/ingest.py b/src/core/ingest.py
index 0c6e02f..fb56393 100644
--- a/src/core/ingest.py
+++ b/src/core/ingest.py
@@ -206,12 +206,13 @@ class DocumentIngestor:
             for i, c in enumerate(chunks)
         ]
 
-        self.collection.add(
-            ids=ids,
-            embeddings=embeddings,
-            documents=texts,
-            metadatas=metadatas,
-        )
+        with self.db.write_lock:
+            self.collection.add(
+                ids=ids,
+                embeddings=embeddings,
+                documents=texts,
+                metadatas=metadatas,
+            )
 
         return len(chunks)
 
@@ -226,10 +227,11 @@ class DocumentIngestor:
     def _remove_by_source(self, file_name: str) -> None:
         """按 source_file 删除已有 chunks."""
         try:
-            existing = self.collection.get(
-                where={"source_file": file_name}
-            )
-            if existing and existing["ids"]:
-                self.collection.delete(ids=existing["ids"])
+            with self.db.write_lock:
+                existing = self.collection.get(
+                    where={"source_file": file_name}
+                )
+                if existing and existing["ids"]:
+                    self.collection.delete(ids=existing["ids"])
         except Exception:
             pass  # collection 为空时 get 可能抛异常
diff --git a/src/core/search.py b/src/core/search.py
index 48e883d..5d84f1d 100644
--- a/src/core/search.py
+++ b/src/core/search.py
@@ -75,14 +75,15 @@ class Searcher:
         return sorted(sources)
 
     def delete_by_source(self, file_name: str) -> bool:
-        """按文件名删除文档."""
+        """按文件名删除文档 (线程安全)."""
         try:
-            existing = self.collection.get(
-                where={"source_file": file_name}
-            )
-            if existing and existing["ids"]:
-                self.collection.delete(ids=existing["ids"])
-                return True
+            with self.db.write_lock:
+                existing = self.collection.get(
+                    where={"source_file": file_name}
+                )
+                if existing and existing["ids"]:
+                    self.collection.delete(ids=existing["ids"])
+                    return True
         except Exception:
             pass
         return False