diff --git a/reme/components/file_chunker/base_file_chunker.py b/reme/components/file_chunker/base_file_chunker.py index 5f4c9081..a4ad0d57 100644 --- a/reme/components/file_chunker/base_file_chunker.py +++ b/reme/components/file_chunker/base_file_chunker.py @@ -3,6 +3,7 @@ import asyncio import contextvars import os +import threading from abc import abstractmethod from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -42,26 +43,45 @@ class BaseFileChunker(BaseComponent): super().__init__(**kwargs) self.supported_extensions: list[str] = supported_extensions or [] self._worker: ThreadPoolExecutor | None = None - self._worker_lock = asyncio.Lock() + self._worker_lock = threading.Lock() + self._worker_state_lock = threading.Lock() self._worker_generation = 0 self._worker_closing = False async def _start(self) -> None: - self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="reme-chunker") - self._worker_generation += 1 - self._worker_closing = False + await self._acquire_worker_lock() + try: + with self._worker_state_lock: + self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="reme-chunker") + self._worker_generation += 1 + self._worker_closing = False + finally: + self._worker_lock.release() async def _close(self) -> None: - self._worker_closing = True - self._worker_generation += 1 + with self._worker_state_lock: + self._worker_closing = True + self._worker_generation += 1 await _finish_work(asyncio.create_task(self._close_worker())) async def _close_worker(self) -> None: - async with self._worker_lock: - if self._worker is not None: - # The admission lock is held through completion, so shutdown has no running work. - self._worker.shutdown(wait=True) + await self._acquire_worker_lock() + try: + with self._worker_state_lock: + self._worker_closing = True + self._worker_generation += 1 + worker = self._worker self._worker = None + if worker is not None: + worker.shutdown(wait=True) + finally: + self._worker_lock.release() + + async def _acquire_worker_lock(self) -> None: + """Acquire cross-loop admission without blocking an event loop thread.""" + # Ownership spans awaited work, so a synchronous lock context would block this loop. + while not self._worker_lock.acquire(blocking=False): # pylint: disable=consider-using-with + await asyncio.sleep(0.01) async def _chunk_in_worker( self, @@ -69,35 +89,50 @@ class BaseFileChunker(BaseComponent): chunk: Callable[[Path, str], tuple[FileNode, list[FileChunk]]], ) -> tuple[FileNode, list[FileChunk]]: """Admit one parse at a time without accumulating an executor work queue.""" - generation = self._worker_generation + with self._worker_state_lock: + generation = self._worker_generation file_path = Path(path) rel_path = self.to_workspace_relative(path) - async with self._worker_lock: - if self._worker_closing or generation != self._worker_generation: - raise RuntimeError("File chunker closed while waiting to parse") - worker = self._worker - temporary_worker = worker is None - if worker is None: - # Standalone chunk() remains usable without start(), with no persistent threads. - worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="reme-chunker") - try: - context = contextvars.copy_context() - future = asyncio.get_running_loop().run_in_executor( - worker, - context.run, - self._chunk_unchanged_file, - file_path, - rel_path, - chunk, - ) - result, version = await _finish_work(future) + await self._acquire_worker_lock() + temporary_worker = False + try: + with self._worker_state_lock: if self._worker_closing or generation != self._worker_generation: - raise RuntimeError("File chunker closed while parsing") - check_file_version(file_path, version, file_path.stat()) - return result - finally: - if temporary_worker: - worker.shutdown(wait=True) + raise RuntimeError("File chunker closed while waiting to parse") + worker = self._worker + temporary_worker = worker is None + if temporary_worker: + # Standalone chunk() remains usable without start(), with no persistent threads. + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="reme-chunker") as worker: + return await self._run_chunk(worker, generation, file_path, rel_path, chunk) + assert worker is not None + return await self._run_chunk(worker, generation, file_path, rel_path, chunk) + finally: + self._worker_lock.release() + + async def _run_chunk( + self, + worker: ThreadPoolExecutor, + generation: int, + file_path: Path, + rel_path: str, + chunk: Callable[[Path, str], tuple[FileNode, list[FileChunk]]], + ) -> tuple[FileNode, list[FileChunk]]: + context = contextvars.copy_context() + future = asyncio.get_running_loop().run_in_executor( + worker, + context.run, + self._chunk_unchanged_file, + file_path, + rel_path, + chunk, + ) + result, version = await _finish_work(future) + with self._worker_state_lock: + if self._worker_closing or generation != self._worker_generation: + raise RuntimeError("File chunker closed while parsing") + check_file_version(file_path, version, file_path.stat()) + return result @staticmethod def _chunk_unchanged_file( diff --git a/reme/steps/index/init_changes.py b/reme/steps/index/init_changes.py index 40711637..fe52e25e 100644 --- a/reme/steps/index/init_changes.py +++ b/reme/steps/index/init_changes.py @@ -73,7 +73,11 @@ 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) + await self.dispatch_steps( + self.dispatch_step_specs, + changes=changes, + retry_source_changes_before_watch=True, + ) 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 26492e5a..cdb904eb 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -18,6 +18,9 @@ from ...enumeration import ComponentEnum from ...schema import FileChunk, FileNode from ...utils.file_utils import FileChangedDuringChunkingError, check_file_version +_SOURCE_CHANGE_RETRY_MARKER = "_source_changed" +_INITIAL_SCAN_RECOVERY_PASSES = 1 + class ChangeApplyStep(BaseStep): """Shared added/modified/deleted handling for index update targets.""" @@ -91,7 +94,40 @@ class ChangeApplyStep(BaseStep): path_exists=lambda p: self._to_abs_path(p).is_file(), path_key=self._change_path_key, ) - results = await self._apply_existing(buckets) + results: list[dict] = [] + pending_buckets = buckets + recover_source_changes = bool(self.context.get("retry_source_changes_before_watch")) + for recovery_pass in range(_INITIAL_SCAN_RECOVERY_PASSES + 1): + batch_results = await self._apply_existing(pending_buckets) + retryable: list[dict] = [] + for result in batch_results: + if result.get(_SOURCE_CHANGE_RETRY_MARKER, False): + retryable.append(result) + if not retryable: + for result in batch_results: + result.pop(_SOURCE_CHANGE_RETRY_MARKER, None) + results.extend(batch_results) + break + if not recover_source_changes: + for result in batch_results: + result.pop(_SOURCE_CHANGE_RETRY_MARKER, None) + results.extend(batch_results) + break + results.extend(result for result in batch_results if not result.pop(_SOURCE_CHANGE_RETRY_MARKER, False)) + pending_buckets = bucket_changes( + [{"change": result["change"], "path": result["path"]} for result in retryable], + path_exists=lambda p: self._to_abs_path(p).is_file(), + path_key=self._change_path_key, + ) + results.extend(await self._apply_deleted(pending_buckets[Change.deleted])) + if recovery_pass == _INITIAL_SCAN_RECOVERY_PASSES: + pending_paths = pending_buckets[Change.added] + pending_buckets[Change.modified] + if pending_paths: + paths = ", ".join(pending_paths) + raise FileChangedDuringChunkingError( + f"Initial scan could not index stable source files after bounded retries: {paths}", + ) + break results.extend(await self._apply_deleted(buckets[Change.deleted])) if persist and results: await self.dump_target() @@ -174,6 +210,7 @@ class ChangeApplyStep(BaseStep): "path": path, "success": False, "error": "not a file", + _SOURCE_CHANGE_RETRY_MARKER: True, }, ) return None @@ -197,6 +234,18 @@ class ChangeApplyStep(BaseStep): ) return None return abs_path, size_bytes + except FileNotFoundError as e: + self.logger.debug(f"Source disappeared while processing {path}; reconciling its current state") + results.append( + { + "change": change.name, + "path": path, + "success": False, + "error": str(e), + _SOURCE_CHANGE_RETRY_MARKER: True, + }, + ) + return None except Exception as e: self.logger.exception(f"Failed to process {path}") results.append({"change": change.name, "path": path, "success": False, "error": str(e)}) @@ -224,7 +273,30 @@ class ChangeApplyStep(BaseStep): self.logger.debug(f"Source changed while processing {path}; retrying") continue self.logger.exception(f"Failed to process {path}") - results.append({"change": change.name, "path": path, "success": False, "error": str(e)}) + results.append( + { + "change": change.name, + "path": path, + "success": False, + "error": str(e), + _SOURCE_CHANGE_RETRY_MARKER: True, + }, + ) + return None + except FileNotFoundError as e: + if attempt < self.source_change_retries: + self.logger.debug(f"Source disappeared while processing {path}; retrying") + continue + self.logger.debug(f"Source disappeared while processing {path}; reconciling its current state") + results.append( + { + "change": change.name, + "path": path, + "success": False, + "error": str(e), + _SOURCE_CHANGE_RETRY_MARKER: True, + }, + ) return None except Exception as e: self.logger.exception(f"Failed to process {path}") @@ -411,6 +483,7 @@ class UpdateIndexStep(ChangeApplyStep): async with self.file_store._maintenance_guard(): # pylint: disable=protected-access current_items: list[tuple[FileNode, list[FileChunk]]] = [] current_paths: list[str] = [] + item: tuple[FileNode, list[FileChunk]] | None = None for item, path in zip(items, ok_paths): node = item[0] try: @@ -428,6 +501,18 @@ class UpdateIndexStep(ChangeApplyStep): "path": path, "success": False, "error": str(exc), + _SOURCE_CHANGE_RETRY_MARKER: True, + } + except FileNotFoundError as exc: + if retry_source_changes: + retry_paths.append(path) + else: + result_by_path[path] = { + "change": change.name, + "path": path, + "success": False, + "error": str(exc), + _SOURCE_CHANGE_RETRY_MARKER: True, } except (OSError, RuntimeError) as exc: result_by_path[path] = { @@ -448,6 +533,10 @@ class UpdateIndexStep(ChangeApplyStep): for result in await super()._try_upsert(change, current_items, current_paths) }, ) + if item is not None: + del item + items.clear() + current_items.clear() retry_results: list[dict] = [] for path in retry_paths: diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index ef0bb30a..126dc8f7 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -54,6 +54,7 @@ from reme.steps.index import ( ) from reme.steps.index._change_batch import bucket_changes from reme.steps.index._watch_rules import WatchRule, build_watch_rules, collect_existing, match_file +from reme.utils.file_utils import FileChangedDuringChunkingError warnings.filterwarnings("ignore", category=DeprecationWarning, module="jieba") warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources") @@ -1014,6 +1015,83 @@ def test_update_index_handles_file_removed_before_stat(): asyncio.run(run()) +def test_update_index_initial_scan_removes_file_disappeared_during_parse(): + """The initial scan removes stale indexed content when a changed source vanishes.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + source = write_file(cwd / "daily" / "a.md", "old contents") + store = LocalFileStore(name="default", embedding_store="") + chunker = DefaultFileChunker(supported_extensions=["md"]) + await store.start() + await chunker.start() + try: + app_ctx = _make_app_context(cwd) + app_ctx.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}} + await store.upsert([await chunker.chunk(source)]) + source.write_text("changed before initial scan", encoding="utf-8") + step = UpdateIndexStep(file_store=store, persist=False, app_context=app_ctx) + + async def disappear_during_parse(path: Path): + path.unlink() + raise FileNotFoundError("file disappeared during parse") + + with patch.object(step, "build_item", side_effect=disappear_during_parse): + response = await step( + RuntimeContext( + changes=[{"change": "modified", "path": str(source)}], + retry_source_changes_before_watch=True, + ), + ) + + assert response.success is True + assert response.answer == [{"change": "deleted", "path": str(source), "success": True}] + assert await store.get_nodes() == [] + finally: + await chunker.close() + await store.close() + + asyncio.run(run()) + + +def test_update_index_initial_scan_fails_when_source_never_stabilizes(): + """A watcher cannot start from an initial scan that exhausted source retries.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + source = write_file(cwd / "daily" / "a.md", "changing source") + store = LocalFileStore(name="default", embedding_store="") + chunker = DefaultFileChunker(supported_extensions=["md"]) + await store.start() + await chunker.start() + try: + app_ctx = _make_app_context(cwd) + app_ctx.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}} + step = UpdateIndexStep(file_store=store, persist=False, app_context=app_ctx) + + async def source_keeps_changing(path: Path): + raise FileChangedDuringChunkingError("source keeps changing") + + with patch.object(step, "build_item", side_effect=source_keeps_changing) as build_item: + with pytest.raises(FileChangedDuringChunkingError, match="after bounded retries"): + await step( + RuntimeContext( + changes=[{"change": "added", "path": str(source)}], + retry_source_changes_before_watch=True, + ), + ) + + assert build_item.call_count == 4 + assert await store.get_nodes() == [] + finally: + await chunker.close() + await store.close() + + asyncio.run(run()) + + def test_update_index_reports_memory_estimation_failure_without_aborting(): """A missing chunker is isolated to the affected file and returned as a normal failure.""" @@ -1108,10 +1186,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_paused = threading.Event() - resume_parse = threading.Event() + parse_started = [threading.Event(), threading.Event()] + resume_parse = [threading.Event(), threading.Event()] watcher_started = asyncio.Event() chunk_calls: list[Path] = [] + watcher_contents: list[str] = [] store = LocalFileStore(name="default", embedding_store="") app_ctx = _make_app_context(cwd) @@ -1147,10 +1226,11 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc def count_chunk(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: result = original_chunk(file_path, rel_path) chunk_calls.append(file_path) - if len(chunk_calls) == 1: - parse_paused.set() - if not resume_parse.wait(timeout=5): - raise TimeoutError("initial parse was not resumed") + attempt = len(chunk_calls) - 1 + if attempt < len(parse_started): + parse_started[attempt].set() + if not resume_parse[attempt].wait(timeout=5): + raise TimeoutError(f"initial parse {attempt + 1} was not resumed") return result from reme.steps.index import watch_changes as watch_changes_module @@ -1167,6 +1247,10 @@ 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]: + 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, @@ -1187,19 +1271,24 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc await store.start() await chunker.start() await job.start() - assert await asyncio.to_thread(parse_paused.wait, 5) + assert await asyncio.to_thread(parse_started[0].wait, 5) + path.write_text("contents after the first parse", encoding="utf-8") + resume_parse[0].set() + assert await asyncio.to_thread(parse_started[1].wait, 5) latest = "latest contents after the initial scan" path.write_text(latest, encoding="utf-8") - resume_parse.set() + resume_parse[1].set() await asyncio.wait_for(watcher_started.wait(), timeout=5) nodes = await store.get_nodes() assert len(nodes) == 1 chunks = [store.file_chunks[chunk_id] for chunk_id in nodes[0].chunk_ids] assert "".join(chunk.text for chunk in chunks) == latest - assert len(chunk_calls) == 2 + assert watcher_contents == [latest] + assert len(chunk_calls) == 3 finally: - resume_parse.set() + for event in resume_parse: + event.set() await job.close() await chunker.close() await store.close() diff --git a/tests/unit/test_file_chunker_workers.py b/tests/unit/test_file_chunker_workers.py index 47c04e28..be64e117 100644 --- a/tests/unit/test_file_chunker_workers.py +++ b/tests/unit/test_file_chunker_workers.py @@ -2,15 +2,18 @@ import asyncio import contextvars +import gc import json import os import threading +import weakref from pathlib import Path from typing import Callable from unittest.mock import patch import pytest from httpx import ASGITransport, AsyncClient +from watchfiles import Change from reme.application import Application from reme.components.application_context import ApplicationContext @@ -204,6 +207,53 @@ async def test_shutdown_rejects_queued_work_and_drains_despite_cancellation(tmp_ assert (await chunker.chunk(path))[1][0].text == "source" +@pytest.mark.asyncio +async def test_shutdown_from_another_event_loop_waits_for_worker_and_closes_executor(tmp_path: Path) -> None: + """Worker admission and shutdown can cross event loops without stranding the executor.""" + path = tmp_path / "source.txt" + path.write_text("source") + chunker = DefaultFileChunker() + entered, release = threading.Event(), threading.Event() + worker_errors: list[BaseException] = [] + original = chunker._chunk_sync + + def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: + entered.set() + assert release.wait(5) + return original(file_path, rel_path) + + def run_on_worker_loop() -> None: + try: + asyncio.run(chunker.chunk(path)) + except BaseException as exc: + worker_errors.append(exc) + + await chunker.start() + with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse): + worker_thread = threading.Thread(target=run_on_worker_loop) + worker_thread.start() + await _wait_started(entered) + closing = asyncio.create_task(chunker.close()) + try: + for _ in range(100): + if chunker._worker_closing: + break + await asyncio.sleep(0.01) + assert chunker._worker_closing + assert not closing.done() + finally: + release.set() + await closing + await asyncio.to_thread(worker_thread.join, 5) + + assert not worker_thread.is_alive() + assert len(worker_errors) == 1 + assert isinstance(worker_errors[0], RuntimeError) + assert "different event loop" not in str(worker_errors[0]) + assert chunker._worker is None + assert not chunker.is_started + + @pytest.mark.asyncio @pytest.mark.parametrize("change", ["modify", "replace", "delete"]) async def test_source_changes_do_not_return_stale_chunks(tmp_path: Path, change: str) -> None: @@ -525,6 +575,76 @@ async def test_changed_batched_file_is_rebuilt_before_publication_and_valid_file assert not step._source_versions +@pytest.mark.asyncio +async def test_stale_batch_chunks_are_released_before_source_retry( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> 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") + 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 + + 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) + + step = UpdateIndexStep(file_store=store, persist=False, batch_max_files=2, app_context=context) + original_build = step._try_build_item + stale_chunk_released: threading.Event | None = None + + async def capture_and_check_stale_chunk( + change: Change, + path: str, + abs_path: Path, + size_bytes: int, + 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: + 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: + 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): + task = asyncio.create_task( + step( + RuntimeContext( + changes=[ + {"change": "modified", "path": str(first_path)}, + {"change": "modified", "path": str(second_path)}, + ], + ), + ), + ) + await _wait_started(entered) + first_path.write_text("latest first") + release.set() + 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 + + @pytest.mark.asyncio @pytest.mark.parametrize("new_size", [12, 100]) async def test_parse_retry_rechecks_size_limit_and_refreshes_memory_estimate(