From 970b7a3b461b97e15070a6ea77d9c8a9f547a2b3 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Thu, 22 Oct 2026 02:12:54 +0800 Subject: [PATCH] Fix initial scan retry scope and stale batches --- reme/steps/index/init_changes.py | 16 +++++--- reme/steps/index/update_changes.py | 2 +- tests/unit/test_background_steps.py | 50 ++++++++++++++--------- tests/unit/test_file_chunker_workers.py | 54 +++++++++++++------------ 4 files changed, 71 insertions(+), 51 deletions(-) diff --git a/reme/steps/index/init_changes.py b/reme/steps/index/init_changes.py index fe52e25e..365305d3 100644 --- a/reme/steps/index/init_changes.py +++ b/reme/steps/index/init_changes.py @@ -73,11 +73,17 @@ class InitChangesStep(BaseStep): self.context["changes"] = changes if changes: self.logger.info(f"[{self.name}] scan {self.monitor_type}:{self.monitor_name}: {counts}") - await self.dispatch_steps( - self.dispatch_step_specs, - changes=changes, - retry_source_changes_before_watch=True, - ) + retry_key = "retry_source_changes_before_watch" + had_retry_value = retry_key in self.context + previous_retry_value = self.context.get(retry_key) + self.context[retry_key] = True + try: + await self.dispatch_steps(self.dispatch_step_specs, changes=changes) + finally: + if had_retry_value: + self.context[retry_key] = previous_retry_value + else: + del self.context[retry_key] else: self.logger.info(f"[{self.name}] {self.monitor_type}:{self.monitor_name} is up to date") self.context.response.metadata["counts"] = counts diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index cdb904eb..0be08e6c 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -189,13 +189,13 @@ class ChangeApplyStep(BaseStep): memory_budget = self._batch_memory_budget() items.append(item) + del item ok_paths.append(path) estimated_bytes += item_bytes if self._batch_is_full(len(items), estimated_bytes, 0, memory_budget): await self._flush_upsert_batch(change, items, ok_paths, estimated_bytes, results) estimated_bytes = 0 memory_budget = self._batch_memory_budget() - del item if items: await self._flush_upsert_batch(change, items, ok_paths, estimated_bytes, results) return results diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index 126dc8f7..b70acd40 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -1150,6 +1150,7 @@ def test_index_update_loop_init_dispatch_updates_store_across_batches(): first = await init_step(ctx) assert first.metadata["counts"] == {"added": 2, "modified": 0, "deleted": 0} + assert "retry_source_changes_before_watch" not in ctx nodes = {n.path: n for n in await fs.get_nodes()} assert set(nodes) == {"daily/a.md", "digest/report.md"} assert all(nodes[p].chunk_ids for p in nodes) @@ -1161,6 +1162,7 @@ def test_index_update_loop_init_dispatch_updates_store_across_batches(): second = await init_step(ctx) assert second.metadata["counts"] == {"added": 1, "modified": 1, "deleted": 1} + assert "retry_source_changes_before_watch" not in ctx assert {(c["change"], Path(c["path"]).name) for c in ctx["changes"]} == { ("modified", "a.md"), ("deleted", "report.md"), @@ -1186,9 +1188,11 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): cwd = Path.cwd() path = write_file(cwd / "daily" / "source.md", "initial contents") - parse_started = [threading.Event(), threading.Event()] - resume_parse = [threading.Event(), threading.Event()] + parse_started = [threading.Event() for _ in range(5)] + resume_parse = [threading.Event() for _ in range(5)] watcher_started = asyncio.Event() + watcher_update_requested = asyncio.Event() + watcher_update_completed = asyncio.Event() chunk_calls: list[Path] = [] watcher_contents: list[str] = [] @@ -1227,16 +1231,12 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc result = original_chunk(file_path, rel_path) chunk_calls.append(file_path) attempt = len(chunk_calls) - 1 - if attempt < len(parse_started): + if attempt in {0, 1, 3, 4}: parse_started[attempt].set() if not resume_parse[attempt].wait(timeout=5): - raise TimeoutError(f"initial parse {attempt + 1} was not resumed") + raise TimeoutError(f"parse attempt {attempt + 1} was not resumed") return result - from reme.steps.index import watch_changes as watch_changes_module - - original_awatch = watch_changes_module.awatch - async def track_watcher_start( *paths: str | Path, watch_filter: Callable[[Change, str], bool], @@ -1247,24 +1247,20 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc poll_delay_ms: int, stop_event: asyncio.Event, ) -> AsyncGenerator[set[tuple[Change, str]], None]: + del paths, watch_filter, recursive, force_polling, debounce, step, poll_delay_ms nodes = await store.get_nodes() watcher_contents.append( "".join(store.file_chunks[chunk_id].text for node in nodes for chunk_id in node.chunk_ids), ) watcher_started.set() - async for changes in original_awatch( - *paths, - watch_filter=watch_filter, - recursive=recursive, - force_polling=force_polling, - debounce=debounce, - step=step, - poll_delay_ms=poll_delay_ms, - stop_event=stop_event, - ): - yield changes + await watcher_update_requested.wait() + yield {(Change.modified, str(path))} + watcher_update_completed.set() + await stop_event.wait() monkeypatch.setattr(chunker, "_chunk_sync", count_chunk) + from reme.steps.index import watch_changes as watch_changes_module + monkeypatch.setattr(watch_changes_module, "awatch", track_watcher_start) try: @@ -1286,6 +1282,22 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc assert "".join(chunk.text for chunk in chunks) == latest assert watcher_contents == [latest] assert len(chunk_calls) == 3 + + path.write_text("watcher update before first parse", encoding="utf-8") + watcher_update_requested.set() + assert await asyncio.to_thread(parse_started[3].wait, 5) + path.write_text("watcher update before retry parse", encoding="utf-8") + resume_parse[3].set() + assert await asyncio.to_thread(parse_started[4].wait, 5) + path.write_text("latest watcher update source", encoding="utf-8") + resume_parse[4].set() + + await asyncio.wait_for(watcher_update_completed.wait(), timeout=5) + assert len(chunk_calls) == 5 + assert job._task is not None and not job._task.done() + nodes = await store.get_nodes() + chunks = [store.file_chunks[chunk_id] for chunk_id in nodes[0].chunk_ids] + assert "".join(chunk.text for chunk in chunks) == latest finally: for event in resume_parse: event.set() diff --git a/tests/unit/test_file_chunker_workers.py b/tests/unit/test_file_chunker_workers.py index be64e117..1d0a86e7 100644 --- a/tests/unit/test_file_chunker_workers.py +++ b/tests/unit/test_file_chunker_workers.py @@ -7,6 +7,7 @@ import json import os import threading import weakref +from contextlib import asynccontextmanager from pathlib import Path from typing import Callable from unittest.mock import patch @@ -582,27 +583,25 @@ async def test_stale_batch_chunks_are_released_before_source_retry( ) -> None: """The retry parse does not overlap the retained stale chunk batch in memory.""" monkeypatch.chdir(tmp_path) - first_path, second_path = tmp_path / "first.txt", tmp_path / "second.txt" - first_path.write_text("original first") - second_path.write_text("original second") + source_path = tmp_path / "source.txt" + source_path.write_text("original source") context = ApplicationContext(workspace_dir=str(tmp_path)) chunker = DefaultFileChunker(app_context=context) context.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}} async with LocalFileStore(embedding_store="") as store, chunker: - await store.upsert([await chunker.chunk(first_path), await chunker.chunk(second_path)]) - first_path.write_text("parsed first") - second_path.write_text("parsed second") - entered, release = threading.Event(), threading.Event() - original_chunk = chunker.chunk + await store.upsert([await chunker.chunk(source_path)]) + source_path.write_text("parsed source") + flush_waiting = asyncio.Event() + original_guard = store._maintenance_guard - async def pause_second(file_path: Path) -> tuple[FileNode, list[FileChunk]]: - if file_path == second_path: - entered.set() - while not release.is_set(): - await asyncio.sleep(0.01) - return await original_chunk(file_path) + @asynccontextmanager + async def track_store_guard(): + flush_waiting.set() + async with original_guard(): + yield - step = UpdateIndexStep(file_store=store, persist=False, batch_max_files=2, app_context=context) + monkeypatch.setattr(store, "_maintenance_guard", track_store_guard) + step = UpdateIndexStep(file_store=store, persist=False, batch_max_files=1, app_context=context) original_build = step._try_build_item stale_chunk_released: threading.Event | None = None @@ -614,35 +613,38 @@ async def test_stale_batch_chunks_are_released_before_source_retry( results: list[dict], ) -> tuple[tuple[FileNode, list[FileChunk]], int] | None: nonlocal stale_chunk_released - if Path(path) == first_path and stale_chunk_released is not None: + if stale_chunk_released is not None: gc.collect() assert stale_chunk_released.is_set() built = await original_build(change, path, abs_path, size_bytes, results) - if Path(path) == first_path and stale_chunk_released is None and built is not None: + if stale_chunk_released is None and built is not None: stale_chunk_released = threading.Event() weakref.finalize(built[0][1][0], stale_chunk_released.set) return built monkeypatch.setattr(step, "_try_build_item", capture_and_check_stale_chunk) - with patch.object(chunker, "chunk", side_effect=pause_second): + await store._maintenance_lock.acquire() + try: task = asyncio.create_task( step( RuntimeContext( - changes=[ - {"change": "modified", "path": str(first_path)}, - {"change": "modified", "path": str(second_path)}, - ], + changes=[{"change": "modified", "path": str(source_path)}], + retry_source_changes_before_watch=True, ), ), ) - await _wait_started(entered) - first_path.write_text("latest first") - release.set() - response = await task + await asyncio.wait_for(flush_waiting.wait(), timeout=5) + source_path.write_text("latest source with a different size") + finally: + store._maintenance_lock.release() + response = await task assert response.success is True assert stale_chunk_released is not None and stale_chunk_released.is_set() assert not step._source_versions + nodes = await store.get_nodes() + assert len(nodes) == 1 + assert store.file_chunks[nodes[0].chunk_ids[0]].text == "latest source with a different size" @pytest.mark.asyncio