fix: coordinate vector migration with concurrent searches and note saves
This commit is contained in:
@@ -190,9 +190,49 @@ async def search_remote(query: str, *, top_k: int, accept_local=False, strict=Fa
|
|||||||
if batch is None:
|
if batch is None:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
if not await _prepare_for_search([batch], strict):
|
||||||
|
return None
|
||||||
return await asyncio.to_thread(_search_space, batch, top_k, strict)
|
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):
|
def _search_space(batch, top_k, strict):
|
||||||
record_embedding(attempted_space={"model_id": batch.space_id, "dimensions": batch.dimensions})
|
record_embedding(attempted_space={"model_id": batch.space_id, "dimensions": batch.dimensions})
|
||||||
try:
|
try:
|
||||||
@@ -209,7 +249,6 @@ def _search_space(batch, top_k, strict):
|
|||||||
if strict:
|
if strict:
|
||||||
raise ValueError("semantic index missing")
|
raise ValueError("semantic index missing")
|
||||||
return None
|
return None
|
||||||
_ensure_table(conn)
|
|
||||||
result = space_index.search(conn, batch, top_k)
|
result = space_index.search(conn, batch, top_k)
|
||||||
record_embedding(source=batch.source, model_id=batch.space_id,
|
record_embedding(source=batch.source, model_id=batch.space_id,
|
||||||
dimensions=batch.dimensions, fallback_reason=None)
|
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:
|
if batch is None:
|
||||||
return None
|
return None
|
||||||
batches[policy] = batch
|
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)
|
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")
|
raise ValueError("embedding policies changed while querying")
|
||||||
ranked = []
|
ranked = []
|
||||||
for policy, batch in batches.items():
|
for policy, batch in batches.items():
|
||||||
_ensure_table(conn)
|
|
||||||
ranked.append(space_index.search(conn, batch, top_k, policy))
|
ranked.append(space_index.search(conn, batch, top_k, policy))
|
||||||
spaces = [{"source": b.source, "model_id": b.space_id, "dimensions": b.dimensions,
|
spaces = [{"source": b.source, "model_id": b.space_id, "dimensions": b.dimensions,
|
||||||
"local_only": policy} for policy, b in batches.items()]
|
"local_only": policy} for policy, b in batches.items()]
|
||||||
|
|||||||
@@ -1,12 +1,42 @@
|
|||||||
"""Persistent vec0 indexes derived from durable routed vectors, one per space/dimension."""
|
"""Persistent vec0 indexes derived from durable routed vectors, one per space/dimension."""
|
||||||
import hashlib
|
import hashlib
|
||||||
import json
|
import json
|
||||||
|
import threading
|
||||||
|
|
||||||
import sqlite_vec
|
import sqlite_vec
|
||||||
|
|
||||||
from app.retrieval.vectorstore import VectorHit
|
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):
|
def table_name(space, dimensions):
|
||||||
return 'routed_vec_' + hashlib.sha256(json.dumps([space, dimensions]).encode()).hexdigest()
|
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):
|
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.
|
# 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=?'
|
where = '' if policy is None else ' AND b.embedding_local_only=?'
|
||||||
params = () if policy is None else (int(policy),)
|
params = () if policy is None else (int(policy),)
|
||||||
|
|||||||
@@ -1,7 +1,14 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
from functools import wraps
|
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):
|
def serialized_vault_mutation(operation):
|
||||||
@@ -9,7 +16,7 @@ def serialized_vault_mutation(operation):
|
|||||||
|
|
||||||
@wraps(operation)
|
@wraps(operation)
|
||||||
async def wrapped(*args, **kwargs):
|
async def wrapped(*args, **kwargs):
|
||||||
async with _vault_mutation_lock:
|
async with vault_mutation_lock():
|
||||||
return await operation(*args, **kwargs)
|
return await operation(*args, **kwargs)
|
||||||
|
|
||||||
return wrapped
|
return wrapped
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ from app.errors import ApiError
|
|||||||
from app.knowledge.parser import parse_note
|
from app.knowledge.parser import parse_note
|
||||||
from app.services.note_service import index_note, prepare_note_index
|
from app.services.note_service import index_note, prepare_note_index
|
||||||
from app.database.db import connect, transaction
|
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.retrieval.vectorstore import SqliteVecStore
|
||||||
from app.local_models.runtime import LocalEmbedding
|
from app.local_models.runtime import LocalEmbedding
|
||||||
from app.services import note_service
|
from app.services import note_service
|
||||||
@@ -116,7 +116,7 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob:
|
|||||||
prepared_notes.append((parsed, prepared))
|
prepared_notes.append((parsed, prepared))
|
||||||
# All network/model awaits precede the transaction. The concrete SQLite
|
# All network/model awaits precede the transaction. The concrete SQLite
|
||||||
# methods below complete synchronously despite their async interfaces.
|
# 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()}:
|
if _scan_vault() != docs or saved_records != {key: repository.get_note_record(key) for key in _pending_notes()}:
|
||||||
raise ApiError(409, "INDEX_SNAPSHOT_CHANGED", "笔记在计算期间发生变化,稍后重新计算。")
|
raise ApiError(409, "INDEX_SNAPSHOT_CHANGED", "笔记在计算期间发生变化,稍后重新计算。")
|
||||||
conn = connect()
|
conn = connect()
|
||||||
@@ -263,7 +263,7 @@ async def _refresh_saved_note(note_id: str) -> None:
|
|||||||
prepared = await prepare_note_index(parsed, strict=True)
|
prepared = await prepare_note_index(parsed, strict=True)
|
||||||
if isinstance(note_service.embedding, LocalEmbedding) and parsed.blocks and prepared[1] is None:
|
if isinstance(note_service.embedding, LocalEmbedding) and parsed.blocks and prepared[1] is None:
|
||||||
raise ApiError(503, "EMBEDDING_UNAVAILABLE", "笔记已保存,后台向量计算未完成。")
|
raise ApiError(503, "EMBEDDING_UNAVAILABLE", "笔记已保存,后台向量计算未完成。")
|
||||||
async with _vault_mutation_lock:
|
async with vault_mutation_lock():
|
||||||
current = repository.get_note_record(note_id)
|
current = repository.get_note_record(note_id)
|
||||||
if current != record or note_service._read_markdown(record.file_path) != markdown:
|
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.
|
# Another save or rename won the race; leave the durable queue entry intact.
|
||||||
|
|||||||
@@ -112,6 +112,115 @@ def test_legacy_vectors_migrate_without_document_embedding(runtime):
|
|||||||
asyncio.run(scenario())
|
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"])
|
@pytest.mark.parametrize("outcome", ["api", "api_failure", "missing_space"])
|
||||||
def test_benchmark_reports_actual_embedding_and_fallback(runtime, outcome):
|
def test_benchmark_reports_actual_embedding_and_fallback(runtime, outcome):
|
||||||
from app.benchmarks import service
|
from app.benchmarks import service
|
||||||
|
|||||||
@@ -10,6 +10,10 @@
|
|||||||
|
|
||||||
数据库搜索及旧空间转换在线程中运行,不阻塞异步服务事件循环。首次转换仍会占用 SQLite 写事务;超大库应进一步评估迁移耗时。新增空间保留原空间索引,暂不自动清理历史空间。
|
数据库搜索及旧空间转换在线程中运行,不阻塞异步服务事件循环。首次转换仍会占用 SQLite 写事务;超大库应进一步评估迁移耗时。新增空间保留原空间索引,暂不自动清理历史空间。
|
||||||
|
|
||||||
|
首次转换使用进程内迁移锁串行执行,在打开检索读快照之前以 `BEGIN IMMEDIATE` 获取写事务;等待迁移时不持有读事务,避免多个搜索从读锁升级写锁发生冲突。迁移完成后二次检查即可复用。普通检索和隐私分区检索共用该流程,已有索引的检索仅执行读取。
|
||||||
|
|
||||||
|
首次转换同时通过 Vault 的异步写入锁与笔记保存、后台索引协调,避免保存事务读取旧 Block 后与迁移争抢写锁。等待是异步的;取消检索时,仍等待迁移线程结束后才释放锁。已就绪的空间直接走只读检查,不进入写入队列。写入锁按事件循环生命周期创建,避免重启后复用已关闭循环的锁。
|
||||||
|
|
||||||
## 外部新增文件
|
## 外部新增文件
|
||||||
|
|
||||||
文件树检测到新增 Markdown 后先登记元数据与全文索引,然后为每条新笔记持久化 `note_vectors_pending:<note_id>`,逐笔记在后台计算。不再因新增文件设置全库重建标记;已存在的全库待处理标记仍保留,以免跳过之前未完成的任务。
|
文件树检测到新增 Markdown 后先登记元数据与全文索引,然后为每条新笔记持久化 `note_vectors_pending:<note_id>`,逐笔记在后台计算。不再因新增文件设置全库重建标记;已存在的全库待处理标记仍保留,以免跳过之前未完成的任务。
|
||||||
@@ -20,5 +24,7 @@
|
|||||||
|
|
||||||
- 隔离数据库运行后端全套测试:637 项通过;最终覆盖检查优化另运行检索及后台工作区回归,96 项通过。
|
- 隔离数据库运行后端全套测试:637 项通过;最终覆盖检查优化另运行检索及后台工作区回归,96 项通过。
|
||||||
- 新增用例验证同模型不同维度并存、重新连接复用且查询不解码向量 JSON、旧表迁移只计算查询向量,以及外部新增文件不触发全量重建。
|
- 新增用例验证同模型不同维度并存、重新连接复用且查询不解码向量 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 耗时,不代表端到端加速比例。
|
- `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。
|
- 原始结果:[2026-09-06-sqlite-vec-search.json](performance/2026-09-06-sqlite-vec-search.json)。基准可用后端虚拟环境 Python 直接执行上述脚本,不读写用户 Vault。
|
||||||
|
|||||||
Reference in New Issue
Block a user