Fix snapshot backfill orchestration regressions (#1854)

* fix snapshot backfill orchestration regressions

* address snapshot orchestration review

* fix admin process projection fixture

* fix noresult projection fixture

archivebox/cli/archivebox_extract.py

@@ -7,6 +7,7 @@ __command__ = "archivebox extract"
 
 import sys
 from collections import defaultdict
+from contextlib import redirect_stdout
 
 import rich_click as click
 
@@ -81,6 +82,12 @@ def _run_snapshot_requests(requested: dict[str, set[str]], *, wait: bool, show_p
         if snapshot is not None and not plugin_names:
             plugin_names.update(get_enabled_plugins(config=get_config(crawl=snapshot.crawl, snapshot=snapshot)))
 
+    if wait and any("search_backend_sonic" in plugin_names for plugin_names in requested.values()):
+        from archivebox.core.takeover_util import ensure_daemon_stack
+
+        with redirect_stdout(sys.stderr):
+            ensure_daemon_stack(reason="Sonic snapshot indexing")
+
     # Explicit extraction resumes open/paused snapshots at the snapshot level.
     # Sealed snapshots stay sealed during targeted maintenance backfills.
     if wait:

archivebox/services/runner.py

@@ -127,6 +127,12 @@ def _count_selected_hooks(catalog: PluginCatalog, selected_plugins: list[str] |
     return sum(1 for plugin in selected.values() for hook in plugin.hooks if "CrawlSetup" in hook.name or "Snapshot" in hook.name)
 
 
+def _enable_requested_plugins(config: dict[str, Any], plugins: PluginCatalog) -> None:
+    for plugin in plugins.values():
+        if plugin.enabled_key in plugin.config.properties:
+            config[plugin.enabled_key] = True
+
+
 def _is_nonfatal_setup_hook(plugin_name: str, hook_name: str) -> bool:
     return plugin_name == "chrome" and hook_name.endswith("_chrome_kill_zombies")
 
@@ -203,7 +209,7 @@ class CrawlRunner:
         self.interactive_interrupts = interactive_interrupts
         self.config_overrides = dict(config_overrides or {})
 
-        SnapshotService(
+        self.snapshot_service = SnapshotService(
             self.bus,
             crawl_id=str(crawl.id),
         )
@@ -372,7 +378,7 @@ class CrawlRunner:
         await asyncio.gather(*done, *pending, return_exceptions=True)
         self.snapshot_tasks.clear()
 
-    async def wait_for_snapshot_tasks(self) -> None:
+    async def wait_for_snapshot_tasks(self, *, enqueue_projected: bool = True) -> None:
         task_errors: list[Exception] = []
         stop_scheduling = False
         while True:
@@ -399,7 +405,8 @@ class CrawlRunner:
                     raise ExceptionGroup("One or more snapshot tasks failed", task_errors)
                 if stop_scheduling:
                     return
-                await self.enqueue_pending_snapshots_from_projection()
+                if enqueue_projected:
+                    await self.enqueue_pending_snapshots_from_projection()
                 if not self.snapshot_tasks:
                     return
                 continue
@@ -425,7 +432,7 @@ class CrawlRunner:
                 await self.crawl_is_cancelled() or (await self.crawl_is_paused() and not self.allow_maintenance_on_inactive_crawl)
             ):
                 stop_scheduling = True
-            if not stop_scheduling:
+            if not stop_scheduling and enqueue_projected:
                 await self.enqueue_pending_snapshots_from_projection()
 
     async def heartbeat_active_leases(self) -> None:
@@ -444,17 +451,17 @@ class CrawlRunner:
         active_snapshot_ids = [snapshot_id for snapshot_id, task in self.snapshot_tasks.items() if not task.done()]
 
         from archivebox.crawls.models import Crawl
-        from archivebox.core.models import Snapshot
 
         await Crawl.objects.filter(id=self.crawl.id, status=Crawl.StatusChoices.STARTED).aupdate(
             retry_at=lease_until,
             modified_at=timezone.now(),
         )
-        if active_snapshot_ids:
-            await Snapshot.objects.filter(id__in=active_snapshot_ids, status=Snapshot.StatusChoices.STARTED).aupdate(
-                retry_at=lease_until,
-                modified_at=timezone.now(),
-            )
+        for snapshot_id in active_snapshot_ids:
+            renewed = await self.snapshot_service.renew_lease(snapshot_id, lease_until)
+            if renewed is False:
+                task = self.snapshot_tasks.get(snapshot_id)
+                if task is not None and not task.done():
+                    task.cancel()
 
     async def drain_snapshot_tasks(self) -> None:
         task_errors: list[Exception] = []
@@ -783,6 +790,8 @@ class CrawlRunner:
         output_dir = Path(self.crawl_output_dir)
         plugins = self.catalog.select(self.selected_plugins)
         config["ABX_RUNTIME"] = "archivebox"
+        if self.requested_plugins is not None:
+            _enable_requested_plugins(config, plugins)
         setup_hooks = plugins.hooks("CrawlSetup")
         abx_snapshot = AbxSnapshot(
             id=snapshot["id"],
@@ -842,15 +851,27 @@ class CrawlRunner:
         async def on_archivebox_CrawlStartEvent(event: CrawlStartEvent) -> None:
             if event.event_id != self.root_crawl_start_event_id:
                 return
-            for snapshot_id in snapshot_ids:
-                if sum(1 for task in self.snapshot_tasks.values() if not task.done()) >= self.max_concurrent_snapshots:
-                    break
-                if await self.crawl_is_cancelled():
-                    break
-                if await self.crawl_is_paused() and not self.allow_maintenance_on_inactive_crawl:
-                    break
-                await self.enqueue_snapshot(snapshot_id)
-            await self.wait_for_snapshot_tasks()
+            if self.initial_snapshot_ids is not None:
+                remaining_snapshot_ids = iter(snapshot_ids)
+                while True:
+                    batch = list(zip(range(self.max_concurrent_snapshots), remaining_snapshot_ids, strict=False))
+                    if not batch:
+                        return
+                    for _slot, snapshot_id in batch:
+                        if await self.crawl_is_cancelled():
+                            return
+                        if await self.crawl_is_paused() and not self.allow_maintenance_on_inactive_crawl:
+                            return
+                        await self.enqueue_snapshot(snapshot_id)
+                    await self.wait_for_snapshot_tasks(enqueue_projected=False)
+            else:
+                for snapshot_id in snapshot_ids[: self.max_concurrent_snapshots]:
+                    if await self.crawl_is_cancelled():
+                        break
+                    if await self.crawl_is_paused():
+                        break
+                    await self.enqueue_snapshot(snapshot_id)
+                await self.wait_for_snapshot_tasks()
 
         async def on_archivebox_CrawlEvent(event: CrawlEvent) -> None:
             if event.event_id != self.root_crawl_event_id:
@@ -1001,6 +1022,8 @@ class CrawlRunner:
             derived_config = normalize_runtime_config(self.derived_config)
             output_dir = Path(snapshot["output_dir"])
             plugins = self.catalog.select(snapshot_selected_plugins) if snapshot_selected_plugins else self.catalog
+            if self.requested_plugins is not None:
+                _enable_requested_plugins(config, plugins)
             abx_snapshot = AbxSnapshot(
                 id=snapshot["id"],
                 url=snapshot["url"],

archivebox/services/snapshot_service.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from pathlib import Path
 
 from asgiref.sync import sync_to_async
@@ -128,6 +129,7 @@ class SnapshotService(BaseService):
     def __init__(self, bus, *, crawl_id: str):
         self.crawl_id = crawl_id
         self._run_ownership: dict[str, tuple[str, object, bool, list[str]]] = {}
+        self._ownership_lock = asyncio.Lock()
         super().__init__(bus)
         self.bus.on(SnapshotEvent, self.on_SnapshotEvent)
         self.bus.on(SnapshotCompletedEvent, self.on_SnapshotCompletedEvent)
@@ -149,14 +151,36 @@ class SnapshotService(BaseService):
             if snapshot.status == Snapshot.StatusChoices.STARTED:
                 await sync_to_async(snapshot.ensure_crawl_symlink, thread_sensitive=True)()
             retry_plugins = [str(name).strip() for name in (snapshot.config or {}).get("RETRY_PLUGINS", []) if str(name).strip()]
-            self._run_ownership[str(event.snapshot_id)] = (str(event.event_id), snapshot.retry_at, was_sealed, retry_plugins)
+            async with self._ownership_lock:
+                self._run_ownership[str(event.snapshot_id)] = (str(event.event_id), snapshot.retry_at, was_sealed, retry_plugins)
+
+    async def renew_lease(self, snapshot_id: str, lease_until) -> bool | None:
+        from archivebox.core.models import Snapshot
+
+        async with self._ownership_lock:
+            ownership = self._run_ownership.get(str(snapshot_id))
+            if ownership is None or ownership[2]:
+                return None
+            event_id, owned_retry_at, was_sealed, retry_plugins = ownership
+            updated = await Snapshot.objects.filter(
+                id=snapshot_id,
+                status=Snapshot.StatusChoices.STARTED,
+                retry_at=owned_retry_at,
+            ).aupdate(
+                retry_at=lease_until,
+                modified_at=timezone.now(),
+            )
+            if updated:
+                self._run_ownership[str(snapshot_id)] = (event_id, lease_until, was_sealed, retry_plugins)
+            return bool(updated)
 
     async def on_SnapshotCompletedEvent(self, event: SnapshotCompletedEvent) -> None:
         snapshot_id = str(event.snapshot_id)
-        ownership = self._run_ownership.get(snapshot_id)
-        if ownership is None or ownership[0] != str(event.event_parent_id):
-            return
-        self._run_ownership.pop(snapshot_id, None)
+        async with self._ownership_lock:
+            ownership = self._run_ownership.get(snapshot_id)
+            if ownership is None or ownership[0] != str(event.event_parent_id):
+                return
+            self._run_ownership.pop(snapshot_id, None)
         _, owned_retry_at, was_sealed, retry_plugins = ownership
         await sync_to_async(finalize_completed_snapshot, thread_sensitive=True)(
             event.snapshot_id,

archivebox/tests/test_cli_update.py

@@ -24,7 +24,7 @@ def test_update_runs_successfully_on_empty_archive(initialized_archive):
 
     assert result.returncode == 0, output
     assert "Phase 1: Draining old archive/ directories" in output
-    assert "Phase 2: Processing all database snapshots" in output
+    assert "Phase 2: Selecting database snapshots with stale filesystem versions" in output
     assert "Updated DB rows:  0" in output
     assert "Sealed crawls:    0" in output
 
@@ -52,7 +52,7 @@ def test_update_reconciles_existing_snapshots(initialized_archive):
 
     output = result.stdout + result.stderr
     assert result.returncode == 0, output
-    assert "Phase 2: Processing all database snapshots" in output
+    assert "Phase 2: Selecting database snapshots with stale filesystem versions" in output
     assert "Updated DB rows:" in output
 
     with use_archivebox_db(initialized_archive):
@@ -147,7 +147,7 @@ def test_update_seals_migrated_snapshots(initialized_archive):
 
     output = result.stdout + result.stderr
     assert result.returncode == 0, output
-    assert "Phase 2: Processing all database snapshots" in output
+    assert "Phase 2: Selecting database snapshots with stale filesystem versions" in output
     assert "Reindexing" not in output
 
     # Check that snapshot remains archived instead of being queued for a full re-crawl.

archivebox/tests/test_takeover_util.py

@@ -3,7 +3,6 @@
 
 import os
 import json
-import re
 import signal
 import subprocess
 import sys
@@ -416,17 +415,10 @@ def test_live_server_keeps_http_runtime_while_update_runs_real_sqlite_indexer(tm
             encoding="utf-8",
             errors="replace",
         )
-        worker_name_match = re.search(r"Worker (worker_runner_update_\d+):", update_stdout)
-        assert worker_name_match, update_stdout
-        wait_for_log(
-            tmp_path / "logs" / f"{worker_name_match.group(1)}.log",
-            "Stopping older ArchiveBox runner process",
-        )
-
-        supervisord_text = wait_for_log_count(supervisord_log, runner_spawn_text, runner_spawn_count + 1, timeout=30)
-        runner_pid_after = int(re.findall(r"spawned: 'worker_runner' with pid (\d+)", supervisord_text)[-1])
-        assert runner_pid_after != runner_pid_before
-        assert pid_is_alive(runner_pid_after)
+        supervisord_text = supervisord_log.read_text(encoding="utf-8", errors="replace")
+        assert supervisord_text.count(runner_spawn_text) == runner_spawn_count
+        assert worker_pid_from_log(server_log, "worker_runner") == runner_pid_before
+        assert pid_is_alive(runner_pid_before)
 
         with use_archivebox_db(tmp_path):
             indexed_results = list(
@@ -439,7 +431,7 @@ def test_live_server_keeps_http_runtime_while_update_runs_real_sqlite_indexer(tm
         server = None
         wait_for_pid_to_disappear(daphne_pid_before, timeout=20)
         wait_for_pid_to_disappear(sonic_pid_before, timeout=20)
-        wait_for_pid_to_disappear(runner_pid_after, timeout=20)
+        wait_for_pid_to_disappear(runner_pid_before, timeout=20)
         assert_no_processes_for_data_dir(tmp_path, timeout=12)
     finally:
         if server is not None:
@@ -479,10 +471,9 @@ def test_live_update_yields_to_server_then_reclaims_real_sqlite_indexing(tmp_pat
         update_supervisor_match = wait_for_log_pattern(update_log, r"Supervisord connected \(pid=(\d+)\)", timeout=90)
         update_supervisor_pid_before = int(update_supervisor_match.group(1))
         update_sonic_pid_before = wait_for_worker_pid_from_log(update_log, "worker_sonic", timeout=90)
-        update_runner_pid_before = wait_for_worker_pid_from_log(update_log, f"worker_runner_update_{update_proc.pid}", timeout=90)
         assert pid_is_alive(update_supervisor_pid_before)
         assert pid_is_alive(update_sonic_pid_before)
-        assert pid_is_alive(update_runner_pid_before)
+        assert "worker_runner_update_" not in update_log.read_text(encoding="utf-8", errors="replace")
         assert "worker_daphne" not in update_log.read_text(encoding="utf-8", errors="replace")
 
         server = start_archivebox_server(tmp_path, port=port, log_name="server-takes-real-sqlite-update.log", env=env)
@@ -507,7 +498,6 @@ def test_live_update_yields_to_server_then_reclaims_real_sqlite_indexing(tmp_pat
         assert pid_is_alive(server_sonic_pid)
         wait_for_pid_to_disappear(update_supervisor_pid_before, timeout=30)
         wait_for_pid_to_disappear(update_sonic_pid_before, timeout=30)
-        wait_for_pid_to_disappear(update_runner_pid_before, timeout=30)
 
         stop_archivebox_process(server, signal.SIGTERM)
         server = None

archivebox/tests/test_ui_admin_archiveresult.py

@@ -17,6 +17,7 @@ pytestmark = pytest.mark.django_db(transaction=True)
 def projected_noresults(snapshot, cached_abxpkg_lib_dir):
     from abx_dl.events import ProcessEvent, SnapshotEvent
     from abx_dl.orchestrator import create_bus
+    from abx_dl.services.archive_result_service import ArchiveResultService as HookArchiveResultService
     from abx_dl.services.process_service import ProcessService as HookProcessService
     from archivebox.core.models import ArchiveResult
     from archivebox.services.archive_result_service import ArchiveResultService
@@ -33,6 +34,7 @@ def projected_noresults(snapshot, cached_abxpkg_lib_dir):
     (staticfile_dir / "input.txt").write_text("plain text without links", encoding="utf-8")
     bus = create_bus(name=f"test_admin_noresults_{snapshot.id}")
     HookProcessService(bus, emit_jsonl=False, interactive_tty=False)
+    HookArchiveResultService(bus, emit_jsonl=False)
     PersistedProcessService(bus)
     ArchiveResultService(bus)
 

archivebox/tests/test_ui_admin_machine.py

@@ -43,6 +43,7 @@ def real_exited_hook_process(tmp_path):
 def real_projected_hash_result(snapshot, cached_abxpkg_lib_dir):
     from abx_dl.events import ProcessEvent, SnapshotEvent
     from abx_dl.orchestrator import create_bus
+    from abx_dl.services.archive_result_service import ArchiveResultService as HookArchiveResultService
     from abx_dl.services.process_service import ProcessService as HookProcessService
     from archivebox.core.models import ArchiveResult
     from archivebox.machine.models import Process
@@ -59,6 +60,7 @@ def real_projected_hash_result(snapshot, cached_abxpkg_lib_dir):
     (snapshot.output_dir / "source.txt").write_text("real admin projection input", encoding="utf-8")
     bus = create_bus(name=f"test_admin_hashes_{snapshot.id}")
     HookProcessService(bus, emit_jsonl=False, interactive_tty=False)
+    HookArchiveResultService(bus, emit_jsonl=False)
     PersistedProcessService(bus)
     ArchiveResultService(bus)
 

archivebox/tests/test_ui_admin_snapshot.py

@@ -399,7 +399,6 @@ class TestSnapshotProgressStats:
     def test_snapshot_admin_progress_uses_expected_hook_total_not_observed_result_count(
         self,
         snapshot,
-        real_hash_projection,
         running_wget_projection,
     ):
         from archivebox.core.admin_site import archivebox_admin
@@ -408,7 +407,6 @@ class TestSnapshotProgressStats:
         from archivebox.config.common import get_config
         from django.urls import resolve
 
-        assert real_hash_projection[1].status == ArchiveResult.StatusChoices.SUCCEEDED
         assert running_wget_projection.status == ArchiveResult.StatusChoices.STARTED
 
         prefetched_snapshot = Snapshot.objects.prefetch_related("archiveresult_set").get(pk=snapshot.pk)
@@ -422,13 +420,13 @@ class TestSnapshotProgressStats:
         stats = admin._get_progress_stats(prefetched_snapshot)
         html = str(admin.status_with_progress(prefetched_snapshot))
 
-        assert expected_total > 2
+        assert expected_total > 1
         assert stats["total"] == expected_total
-        assert stats["succeeded"] == 1
+        assert stats["succeeded"] == 0
         assert stats["running"] == 1
-        assert stats["pending"] == expected_total - 2
-        assert stats["percent"] == int(100 / expected_total)
-        assert f"1/{expected_total} hooks" in html
+        assert stats["pending"] == expected_total - 1
+        assert stats["percent"] == 0
+        assert f"0/{expected_total} hooks" in html
 
     def test_get_progress_stats_sealed(self, snapshot):
         """Test progress stats for sealed snapshot."""