From bab44fb8323977b7d0bbfa0e282e9e4885024676 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Tue, 3 Nov 2026 13:30:57 +0800 Subject: [PATCH] refactor(index): carry source versions with batch items --- .../file_chunker/base_file_chunker.py | 7 ++-- reme/steps/index/update_changes.py | 35 ++++++------------- tests/unit/test_file_chunker_workers.py | 24 +++++++------ 3 files changed, 26 insertions(+), 40 deletions(-) diff --git a/reme/components/file_chunker/base_file_chunker.py b/reme/components/file_chunker/base_file_chunker.py index 8ec75c8d..6d9fed61 100644 --- a/reme/components/file_chunker/base_file_chunker.py +++ b/reme/components/file_chunker/base_file_chunker.py @@ -44,7 +44,7 @@ class BaseFileChunker(BaseComponent): with self._worker_state_lock: self._worker_closing = True self._worker_generation += 1 - await finish_cleanup(asyncio.create_task(self._close_worker())) + await finish_cleanup(self._close_worker()) async def _close_worker(self) -> None: await self._acquire_worker_lock() @@ -76,18 +76,15 @@ class BaseFileChunker(BaseComponent): file_path = Path(path) rel_path = self.to_workspace_relative(path) 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 waiting to parse") worker = self._worker - temporary_worker = worker is None - if temporary_worker: + if worker is None: # 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() diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index 2ff51221..11a74b01 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -407,33 +407,22 @@ class UpdateIndexStep(ChangeApplyStep): reads_file_content = True source_change_retries: int = 1 - def __init__(self, *args: Any, **kwargs: Any) -> None: - super().__init__(*args, **kwargs) - self._source_versions: dict[str, os.stat_result] = {} - def _change_path_key(self, path: str) -> str: """Coalesce absolute and workspace-relative aliases of the same file.""" return os.path.normcase(os.fspath(self._to_abs_path(path).absolute())) - async def _apply_existing(self, buckets: dict[Change, list[str]]) -> list[dict]: - try: - return await super()._apply_existing(buckets) - finally: - # Steps are invocation-scoped; cancelled batches must release their snapshots too. - self._source_versions.clear() - - async def build_item(self, path: Path) -> tuple[FileNode, list[FileChunk]]: + async def build_item(self, path: Path) -> tuple[FileNode, list[FileChunk], os.stat_result]: + """Keep each source version with its parsed item until publication.""" before = path.stat() - item = await self.chunk_file(path) + node, chunks = await self.chunk_file(path) after = path.stat() check_file_version(path, before, after) - self._source_versions[item[0].path] = after - return item + return node, chunks, after async def _try_upsert( self, change: Change, - items: list[tuple[FileNode, list[FileChunk]]], + items: list[tuple[FileNode, list[FileChunk], os.stat_result]], ok_paths: list[str], *, retry_source_changes: bool = True, @@ -442,15 +431,15 @@ class UpdateIndexStep(ChangeApplyStep): retry_paths: list[str] = [] result_by_path: dict[str, dict] = {} async with self.file_store._maintenance_guard(): # pylint: disable=protected-access - current_items: list[tuple[FileNode, list[FileChunk]]] = [] + current_items: list[tuple[FileNode, list[FileChunk], os.stat_result]] = [] current_paths: list[str] = [] - item: tuple[FileNode, list[FileChunk]] | None = None + item: tuple[FileNode, list[FileChunk], os.stat_result] | None = None for item, path in zip(items, ok_paths): node = item[0] try: check_file_version( self._to_abs_path(node.path), - self._source_versions[node.path], + item[2], self._to_abs_path(node.path).stat(), ) except (FileChangedDuringChunkingError, FileNotFoundError) as exc: @@ -474,8 +463,6 @@ class UpdateIndexStep(ChangeApplyStep): else: current_items.append(item) current_paths.append(path) - finally: - self._source_versions.pop(node.path) if current_items: result_by_path.update( { @@ -509,8 +496,8 @@ class UpdateIndexStep(ChangeApplyStep): result_by_path.update({result["path"]: result for result in retry_results}) return [result_by_path[path] for path in ok_paths] - async def upsert_items(self, items: list[tuple[FileNode, list[FileChunk]]]) -> None: - await self.file_store.upsert(items) + async def upsert_items(self, items: list[tuple[FileNode, list[FileChunk], os.stat_result]]) -> None: + await self.file_store.upsert([(node, chunks) for node, chunks, _ in items]) async def delete_paths(self, paths: list[str]) -> None: await self.file_store.delete(paths) @@ -526,7 +513,7 @@ class UpdateIndexStep(ChangeApplyStep): def estimate_item_memory( self, - item: tuple[FileNode, list[FileChunk]], + item: tuple[FileNode, list[FileChunk], os.stat_result], path: Path, size_bytes: int, ) -> int: diff --git a/tests/unit/test_file_chunker_workers.py b/tests/unit/test_file_chunker_workers.py index fb9e2697..f34900c5 100644 --- a/tests/unit/test_file_chunker_workers.py +++ b/tests/unit/test_file_chunker_workers.py @@ -506,7 +506,6 @@ async def test_existing_content_searchable_during_indexing( assert recovered.success is True nodes = await store.get_nodes() assert store.file_chunks[nodes[0].chunk_ids[0]].text == path.read_text() - assert not step._source_versions @pytest.mark.asyncio @@ -576,7 +575,6 @@ async def test_changed_batched_file_is_rebuilt_before_publication_and_valid_file if change == "oversize": assert response.answer[0]["skipped"] is True assert response.answer[0]["reason"] == "file_too_large" - assert not step._source_versions nodes = {node.path: node for node in await store.get_nodes()} if change == "oversize": assert "first.txt" not in nodes @@ -602,7 +600,6 @@ async def test_changed_batched_file_is_rebuilt_before_publication_and_valid_file assert "first.txt" not in nodes else: assert store.file_chunks[nodes["first.txt"].chunk_ids[0]].text == first_path.read_text() - assert not step._source_versions @pytest.mark.asyncio @@ -640,7 +637,7 @@ async def test_stale_batch_chunks_are_released_before_source_retry( abs_path: Path, size_bytes: int, results: list[dict], - ) -> tuple[tuple[FileNode, list[FileChunk]], int] | None: + ) -> tuple[tuple[FileNode, list[FileChunk], os.stat_result], int] | None: nonlocal stale_chunk_released if stale_chunk_released is not None: gc.collect() @@ -670,7 +667,6 @@ async def test_stale_batch_chunks_are_released_before_source_retry( 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" @@ -708,7 +704,11 @@ async def test_parse_retry_rechecks_size_limit_and_refreshes_memory_estimate( estimated_sizes: list[int] = [] original_estimate = step.estimate_item_memory - def record_estimate(item: tuple[FileNode, list[FileChunk]], file_path: Path, size_bytes: int) -> int: + def record_estimate( + item: tuple[FileNode, list[FileChunk], os.stat_result], + file_path: Path, + size_bytes: int, + ) -> int: estimated_sizes.append(size_bytes) return original_estimate(item, file_path, size_bytes) @@ -736,7 +736,6 @@ async def test_parse_retry_rechecks_size_limit_and_refreshes_memory_estimate( assert estimated_sizes == [new_size] nodes = await store.get_nodes() assert store.file_chunks[nodes[0].chunk_ids[0]].text == path.read_text() - assert not step._source_versions @pytest.mark.asyncio @@ -767,7 +766,6 @@ async def test_index_updates_deduplicate_absolute_and_relative_paths( nodes = await store.get_nodes() assert len(nodes) == 1 assert nodes[0].path == "source.txt" - assert not step._source_versions @pytest.mark.asyncio @@ -788,11 +786,13 @@ async def test_version_is_checked_after_waiting_for_store_lock( await store.upsert([await chunker.chunk(path)]) path.write_text("updated source") built = asyncio.Event() + pending_chunks: list[weakref.ReferenceType[FileChunk]] = [] step = UpdateIndexStep(file_store=store, persist=False, app_context=context) original = step.build_item - async def record_build(file_path: Path) -> tuple[FileNode, list[FileChunk]]: + async def record_build(file_path: Path) -> tuple[FileNode, list[FileChunk], os.stat_result]: item = await original(file_path) + pending_chunks.append(weakref.ref(item[1][0])) built.set() return item @@ -810,7 +810,10 @@ async def test_version_is_checked_after_waiting_for_store_lock( if outcome == "change": response = await task assert response.success is True - assert not step._source_versions + if outcome == "cancel": + await asyncio.sleep(0) + gc.collect() + assert pending_chunks and all(ref() is None for ref in pending_chunks) nodes = await store.get_nodes() expected_text = "latest source" if outcome == "change" else "original source" assert store.file_chunks[nodes[0].chunk_ids[0]].text == expected_text @@ -819,4 +822,3 @@ async def test_version_is_checked_after_waiting_for_store_lock( assert recovered.success is True nodes = await store.get_nodes() assert store.file_chunks[nodes[0].chunk_ids[0]].text == path.read_text() - assert not step._source_versions