diff --git a/backend/app/retrieval/routed_vectors.py b/backend/app/retrieval/routed_vectors.py index 8ea10cd..ab959be 100644 --- a/backend/app/retrieval/routed_vectors.py +++ b/backend/app/retrieval/routed_vectors.py @@ -190,9 +190,49 @@ async def search_remote(query: str, *, top_k: int, accept_local=False, strict=Fa if batch is None: return None + if not await _prepare_for_search([batch], strict): + return None return await asyncio.to_thread(_search_space, batch, top_k, strict) +async def _prepare_indexes(batches): + from app.services.coordination import vault_mutation_lock + def prepare(check_only=False): + conn = connect() + try: + if check_only: + return space_index.is_ready(conn, batches) + space_index.prepare(conn, batches) + finally: + conn.close() + if await asyncio.to_thread(prepare, True): + return + # Share the cooperative gate with saves: never block the event loop on a + # SQLite write lock while a migration owns it in another thread. + async with vault_mutation_lock(): + work = asyncio.create_task(asyncio.to_thread(prepare)) + cancelled = False + while not work.done(): + try: + await asyncio.shield(work) + except asyncio.CancelledError: + cancelled = True + work.result() + if cancelled: + raise asyncio.CancelledError + + +async def _prepare_for_search(batches, strict): + try: + await _prepare_indexes(batches) + return True + except Exception as exc: + record_embedding(fallback_reason='REMOTE_INDEX_UNAVAILABLE') + if strict: + raise ApiError(409, 'SEMANTIC_INDEX_UNAVAILABLE', '向量索引准备失败,请检查索引状态。') from exc + return False + + def _search_space(batch, top_k, strict): record_embedding(attempted_space={"model_id": batch.space_id, "dimensions": batch.dimensions}) try: @@ -209,7 +249,6 @@ def _search_space(batch, top_k, strict): if strict: raise ValueError("semantic index missing") return None - _ensure_table(conn) result = space_index.search(conn, batch, top_k) record_embedding(source=batch.source, model_id=batch.space_id, dimensions=batch.dimensions, fallback_reason=None) @@ -234,6 +273,8 @@ async def _search_partitioned(query: str, policies: set[bool], *, top_k: int, st if batch is None: return None batches[policy] = batch + if not await _prepare_for_search(list(batches.values()), strict): + return None return await asyncio.to_thread(_search_partitions, batches, policies, top_k, strict) @@ -247,7 +288,6 @@ def _search_partitions(batches, policies, top_k, strict): raise ValueError("embedding policies changed while querying") ranked = [] for policy, batch in batches.items(): - _ensure_table(conn) ranked.append(space_index.search(conn, batch, top_k, policy)) spaces = [{"source": b.source, "model_id": b.space_id, "dimensions": b.dimensions, "local_only": policy} for policy, b in batches.items()] diff --git a/backend/app/retrieval/space_index.py b/backend/app/retrieval/space_index.py index ddf9e0f..1947c5b 100644 --- a/backend/app/retrieval/space_index.py +++ b/backend/app/retrieval/space_index.py @@ -1,12 +1,42 @@ """Persistent vec0 indexes derived from durable routed vectors, one per space/dimension.""" import hashlib import json +import threading import sqlite_vec from app.retrieval.vectorstore import VectorHit +_migration_lock = threading.Lock() + + +def is_ready(conn, batches): + return all(conn.execute('SELECT 1 FROM sqlite_master WHERE name=?', + (table_name(batch.space_id, batch.dimensions),)).fetchone() for batch in batches) + + +def prepare(conn, batches): + """Finish lazy writes before opening a search snapshot. Warm searches do not write.""" + from app.retrieval.routed_vectors import _ensure_table + batches = list(batches) + if is_ready(conn, batches): + return + # Waiting holds no read transaction, so a concurrent migration can commit. + with _migration_lock: + if is_ready(conn, batches): + return + conn.execute('BEGIN IMMEDIATE') + try: + _ensure_table(conn) + for batch in batches: + ensure(conn, batch.space_id, batch.dimensions) + conn.execute('COMMIT') + except BaseException: + conn.execute('ROLLBACK') + raise + + def table_name(space, dimensions): return 'routed_vec_' + hashlib.sha256(json.dumps([space, dimensions]).encode()).hexdigest() @@ -40,7 +70,7 @@ def upsert(conn, block_ids, batch): def search(conn, batch, top_k, policy=None): - table = ensure(conn, batch.space_id, batch.dimensions) + table = table_name(batch.space_id, batch.dimensions) # Coverage checks stay relational; no JSON decoding or Python dot products on the hot path. where = '' if policy is None else ' AND b.embedding_local_only=?' params = () if policy is None else (int(policy),) diff --git a/backend/app/services/coordination.py b/backend/app/services/coordination.py index b6dd0e2..6452249 100644 --- a/backend/app/services/coordination.py +++ b/backend/app/services/coordination.py @@ -1,7 +1,14 @@ import asyncio from functools import wraps +from weakref import WeakKeyDictionary -_vault_mutation_lock = asyncio.Lock() +_vault_locks = WeakKeyDictionary() + + +def vault_mutation_lock(): + # Service/test lifecycle restarts must not reuse a lock bound to a closed loop. + loop = asyncio.get_running_loop() + return _vault_locks.setdefault(loop, asyncio.Lock()) def serialized_vault_mutation(operation): @@ -9,7 +16,7 @@ def serialized_vault_mutation(operation): @wraps(operation) async def wrapped(*args, **kwargs): - async with _vault_mutation_lock: + async with vault_mutation_lock(): return await operation(*args, **kwargs) return wrapped diff --git a/backend/app/services/index_service.py b/backend/app/services/index_service.py index 2534b3d..db80a59 100644 --- a/backend/app/services/index_service.py +++ b/backend/app/services/index_service.py @@ -17,7 +17,7 @@ from app.errors import ApiError from app.knowledge.parser import parse_note from app.services.note_service import index_note, prepare_note_index from app.database.db import connect, transaction -from app.services.coordination import _vault_mutation_lock +from app.services.coordination import vault_mutation_lock from app.retrieval.vectorstore import SqliteVecStore from app.local_models.runtime import LocalEmbedding from app.services import note_service @@ -116,7 +116,7 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob: prepared_notes.append((parsed, prepared)) # All network/model awaits precede the transaction. The concrete SQLite # methods below complete synchronously despite their async interfaces. - async with _vault_mutation_lock: + async with vault_mutation_lock(): if _scan_vault() != docs or saved_records != {key: repository.get_note_record(key) for key in _pending_notes()}: raise ApiError(409, "INDEX_SNAPSHOT_CHANGED", "笔记在计算期间发生变化,稍后重新计算。") conn = connect() @@ -263,7 +263,7 @@ async def _refresh_saved_note(note_id: str) -> None: prepared = await prepare_note_index(parsed, strict=True) if isinstance(note_service.embedding, LocalEmbedding) and parsed.blocks and prepared[1] is None: raise ApiError(503, "EMBEDDING_UNAVAILABLE", "笔记已保存,后台向量计算未完成。") - async with _vault_mutation_lock: + async with vault_mutation_lock(): current = repository.get_note_record(note_id) if current != record or note_service._read_markdown(record.file_path) != markdown: # Another save or rename won the race; leave the durable queue entry intact. diff --git a/backend/tests/test_routed_retrieval.py b/backend/tests/test_routed_retrieval.py index 282b8ca..d647100 100644 --- a/backend/tests/test_routed_retrieval.py +++ b/backend/tests/test_routed_retrieval.py @@ -112,6 +112,115 @@ def test_legacy_vectors_migrate_without_document_embedding(runtime): asyncio.run(scenario()) +@pytest.mark.parametrize('partitioned', [False, True]) +def test_concurrent_first_search_serializes_migration_and_warm_search_is_read_only(runtime, monkeypatch, partitioned): + import threading + from app.retrieval import space_index + async def scenario(): + await seed() + table = space_index.table_name('space-a', 3) + conn = connect() + try: + with transaction(conn): + conn.execute(f'DROP TRIGGER {table}_delete') + conn.execute(f'DROP TRIGGER {table}_update') + conn.execute(f'DROP TABLE {table}') + finally: + conn.close() + entered, release, second = threading.Event(), threading.Event(), threading.Event() + original_ensure, original_prepare = space_index.ensure, space_index.prepare + calls = [] + def ensure(*args): + calls.append(1) + entered.set() + assert release.wait(5) + return original_ensure(*args) + def prepare(*args): + if entered.is_set(): + second.set() + return original_prepare(*args) + monkeypatch.setattr(space_index, 'ensure', ensure) + monkeypatch.setattr(space_index, 'prepare', prepare) + batch = routed_vectors.RemoteEmbeddings('space-a', 3, [[1., 0., 0.]]) + async def search(): + if entered.is_set(): + second.set() + await routed_vectors._prepare_indexes([batch]) + if partitioned: + return await asyncio.to_thread(routed_vectors._search_partitions, {False: batch}, {False}, 2, True) + return await asyncio.to_thread(routed_vectors._search_space, batch, 2, True) + tasks = [] + try: + tasks.append(asyncio.create_task(search())) + assert await asyncio.to_thread(entered.wait, 5) + tasks.append(asyncio.create_task(search())) + assert await asyncio.to_thread(second.wait, 5) + release.set() + first, other = await asyncio.gather(*tasks) + assert first == other and len(first) == 2 + assert len(calls) == 1 + # Prepared indexes are reusable even with SQLite query_only enforced. + original_connect = routed_vectors.connect + def read_only(): + connection = original_connect() + connection.execute('PRAGMA query_only=ON') + return connection + monkeypatch.setattr(routed_vectors, 'connect', read_only) + assert await search() == first + finally: + release.set() + await asyncio.gather(*tasks, return_exceptions=True) + asyncio.run(scenario()) + + +@pytest.mark.parametrize('cancel_search', [False, True]) +def test_save_waits_for_migration_even_when_search_is_cancelled(runtime, monkeypatch, cancel_search): + import threading + from app.retrieval import space_index + async def scenario(): + apple, _ = await seed() + table = space_index.table_name('space-a', 3) + conn = connect() + try: + with transaction(conn): + conn.execute(f'DROP TRIGGER {table}_delete') + conn.execute(f'DROP TRIGGER {table}_update') + conn.execute(f'DROP TABLE {table}') + finally: + conn.close() + entered, release = threading.Event(), threading.Event() + original = space_index.ensure + def slow(*args): + entered.set() + assert release.wait(5) + return original(*args) + monkeypatch.setattr(space_index, 'ensure', slow) + # Keep the subsequent vector job queued; test saving and its durable marker. + monkeypatch.setattr(index_service, 'schedule_workspace_rebuild', lambda: None) + query = asyncio.create_task(routed_vectors.search_remote('apple orchard', top_k=2, strict=True)) + save = None + try: + assert await asyncio.to_thread(entered.wait, 5) + if cancel_search: + query.cancel() + save = asyncio.create_task(note_service.update_note(apple.note_id, markdown='Saved during migration', defer_vectors=True)) + await asyncio.sleep(0.02) + assert not save.done() + release.set() + saved = await asyncio.wait_for(save, 5) + assert saved.markdown == 'Saved during migration' + assert (await note_service.get_note(apple.note_id)).markdown == saved.markdown + assert repository.get_index_meta()[f'note_vectors_pending:{apple.note_id}'] == '1' + # Query may observe the saved revision's pending index, but saving must succeed. + result = (await asyncio.gather(query, return_exceptions=True))[0] + if cancel_search: + assert isinstance(result, asyncio.CancelledError) + finally: + release.set() + await asyncio.gather(*([query, save] if save else [query]), return_exceptions=True) + asyncio.run(scenario()) + + @pytest.mark.parametrize("outcome", ["api", "api_failure", "missing_space"]) def test_benchmark_reports_actual_embedding_and_fallback(runtime, outcome): from app.benchmarks import service diff --git a/docs/development/模型隔离向量索引与增量登记.md b/docs/development/模型隔离向量索引与增量登记.md index aff56b4..b36339a 100644 --- a/docs/development/模型隔离向量索引与增量登记.md +++ b/docs/development/模型隔离向量索引与增量登记.md @@ -10,6 +10,10 @@ 数据库搜索及旧空间转换在线程中运行,不阻塞异步服务事件循环。首次转换仍会占用 SQLite 写事务;超大库应进一步评估迁移耗时。新增空间保留原空间索引,暂不自动清理历史空间。 +首次转换使用进程内迁移锁串行执行,在打开检索读快照之前以 `BEGIN IMMEDIATE` 获取写事务;等待迁移时不持有读事务,避免多个搜索从读锁升级写锁发生冲突。迁移完成后二次检查即可复用。普通检索和隐私分区检索共用该流程,已有索引的检索仅执行读取。 + +首次转换同时通过 Vault 的异步写入锁与笔记保存、后台索引协调,避免保存事务读取旧 Block 后与迁移争抢写锁。等待是异步的;取消检索时,仍等待迁移线程结束后才释放锁。已就绪的空间直接走只读检查,不进入写入队列。写入锁按事件循环生命周期创建,避免重启后复用已关闭循环的锁。 + ## 外部新增文件 文件树检测到新增 Markdown 后先登记元数据与全文索引,然后为每条新笔记持久化 `note_vectors_pending:`,逐笔记在后台计算。不再因新增文件设置全库重建标记;已存在的全库待处理标记仍保留,以免跳过之前未完成的任务。 @@ -20,5 +24,7 @@ - 隔离数据库运行后端全套测试:637 项通过;最终覆盖检查优化另运行检索及后台工作区回归,96 项通过。 - 新增用例验证同模型不同维度并存、重新连接复用且查询不解码向量 JSON、旧表迁移只计算查询向量,以及外部新增文件不触发全量重建。 +- 并发迁移回归用同步事件暂停第一个请求的转换,同时发起第二个请求;验证普通/隐私分区路径均只迁移一次、两个请求结果一致,并开启 SQLite `query_only` 验证后续检索不会写入。相关检索与后台工作区测试 91 项通过。 +- 并发保存回归暂停迁移后发起实际 `update_note(..., defer_vectors=True)`,验证保存等待、事件循环仍可运行、放行迁移后正文成功落盘并保留新内容的向量待处理标记;另覆盖搜索被取消时不得提前放行保存。 - `backend/scripts/vector-index-benchmark.py` 使用临时数据库、固定随机种子,比较 4000 条 384 维向量的 top-20,旧新路径结果顺序一致。5 次采样中位数:Python 扫描约 1289.64 ms,sqlite-vec 约 19.67 ms;首次转换约 1329.51 ms。该结果只测向量检索,不包含查询 Embedding、重排与 HTTP 耗时,不代表端到端加速比例。 - 原始结果:[2026-09-06-sqlite-vec-search.json](performance/2026-09-06-sqlite-vec-search.json)。基准可用后端虚拟环境 Python 直接执行上述脚本,不读写用户 Vault。