Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 19 additions & 1 deletion apps/backend/app/workers/analysis_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -409,7 +409,7 @@ def _stage_extract(self, ctx: _StageContext) -> None:
try:
deferred_dependencies: list[ProducedExtraction] = []
for produced in pipeline.iter_run(
sources,
self._sources_with_progress(ctx, sources),
check_cancelled=lambda: self._check_heartbeat(ctx),
):
if any(node.node_kind == "dependency" for node in produced.result.nodes):
Expand All @@ -435,6 +435,24 @@ def _stage_extract(self, ctx: _StageContext) -> None:
finally:
ctx.release_writes = False

def _sources_with_progress(
self, ctx: _StageContext, sources: RepositorySourceStream
) -> Iterator[tuple[str, bytes]]:
"""Publish completed-file progress without counting extractor yields as files."""
total = max(1, self._require_record(ctx).file_count)
published = 35
for completed, source in enumerate(sources, start=1):
yield source
# The pipeline has finished this source before it requests the next one.
progress = 35 + min(35, completed * 35 // max(1, total))
if progress > published:
self._check_heartbeat(ctx)
record = self._require_record(ctx)
record.analysis_stage = "extracting-modules"
record.analysis_progress = progress
self._checkpoint(ctx, "extracting-modules", progress)
published = progress

def _stage_seal(self, ctx: _StageContext) -> None:
"""Seal the building snapshot (commits internally, see the module note)."""

Expand Down
22 changes: 22 additions & 0 deletions apps/backend/tests/test_analysis_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -1039,3 +1039,25 @@ def _boom(_ctx):
assert job.snapshot_id is not None
assert session.scalar(select(func.count()).select_from(RiSnapshot)) == 1
assert session.get(RiSnapshot, job.snapshot_id).state == "completed"


def test_extraction_progress_advances_with_completed_files(session_factory, tmp_path, monkeypatch):
with session_factory() as session:
owner = _owner(session)
record = _repository_with_sources(session, owner, tmp_path / "progress-repo")
AnalysisJobService(session, owner.id).submit(record.id)
worker = AnalysisWorker(session_factory, worker_id="progress", lease_seconds=60)
observed = []
checkpoint = worker._checkpoint

def capture(ctx, stage, progress):
checkpoint(ctx, stage, progress)
observed.append((stage, progress))

monkeypatch.setattr(worker, "_checkpoint", capture)
assert worker.run_once()
extraction = [progress for stage, progress in observed if stage == "extracting-modules"]
assert extraction[0] == 35
assert any(35 < progress <= 70 for progress in extraction)
assert extraction == sorted(set(extraction))
assert max(extraction) < 100
Loading