Preserve queued index jobs during reindex

Dockerfile.multistage

@@ -199,7 +199,7 @@ RUN echo "[+] Initializing image collection..." \
 RUN chmod +x "$CODE_DIR"/bin/*.sh \
     && chown -R "$DEFAULT_PUID:$DEFAULT_PGID" "$LIB_DIR" \
     && chmod g+w "$TMP_DIR" "$LIB_DIR" "$PLAYWRIGHT_BROWSERS_PATH" \
-    && GIT_BINARY="$LIB_DIR/env/bin/git" GALLERYDL_BINARY="$LIB_DIR/env/bin/gallery-dl" FORUMDL_BINARY="$LIB_DIR/env/bin/forum-dl" ABXPKG_INSTALL_TIMEOUT=600 ABXPKG_POSTINSTALL_SCRIPTS=True ABXPKG_MIN_RELEASE_AGE=0 TIMEOUT=600 gosu "$ARCHIVEBOX_USER" archivebox install archivewebpage defuddle forumdl gallerydl git istilldontcareaboutcookies liteparse mercury papersdl parse_rss_urls readability search_backend_ripgrep search_backend_sonic 2>&1 | tee -a /VERSION.txt \
+    && GIT_BINARY="$LIB_DIR/env/bin/git" GALLERYDL_BINARY="$LIB_DIR/env/bin/gallery-dl" FORUMDL_BINARY="$LIB_DIR/env/bin/forum-dl" ABXPKG_INSTALL_TIMEOUT=600 ABXPKG_POSTINSTALL_SCRIPTS=True ABXPKG_MIN_RELEASE_AGE=0 TIMEOUT=600 gosu "$ARCHIVEBOX_USER" archivebox install archivewebpage defuddle forumdl gallerydl git istilldontcareaboutcookies liteparse mercury opendataloader papersdl parse_rss_urls readability search_backend_ripgrep search_backend_sonic 2>&1 | tee -a /VERSION.txt \
     && "$LIB_DIR/env/bin/chromium" --version | tee -a /VERSION.txt \
     && "$LIB_DIR/pip/packages/papers-dl/venv/bin/papers-dl" --version | tee -a /VERSION.txt \
     && /usr/bin/rg --version | head -1 | tee -a /VERSION.txt \

archivebox/cli/archivebox_extract.py

@@ -116,6 +116,7 @@ def run_plugins(
     wait: bool = True,
     emit_results: bool = True,
     show_progress: bool = True,
+    preserve_queued: bool = False,
 ) -> int:
     """
     Run plugins on Snapshots from input.
@@ -219,6 +220,18 @@ def run_plugins(
         else:
             requested_rows.add((snapshot_id, plugin_name, ""))
 
+    queued_rows: set[tuple[str, str, str]] = set()
+    if preserve_queued and requested_rows:
+        queued_rows = {
+            (str(snapshot_id), plugin_name, hook_name)
+            for snapshot_id, plugin_name, hook_name in ArchiveResult.objects.filter(
+                snapshot_id__in=existing_snapshot_ids,
+                plugin__in={plugin_name for _snapshot_id, plugin_name, _hook_name in requested_rows},
+                status=ArchiveResult.StatusChoices.QUEUED,
+            ).values_list("snapshot_id", "plugin", "hook_name")
+        }
+    rows_to_queue = requested_rows - queued_rows
+
     reset_fields = {
         "status": ArchiveResult.StatusChoices.QUEUED,
         "output_str": "",
@@ -230,23 +243,34 @@ def run_plugins(
         "end_ts": None,
         "modified_at": timezone.now(),
     }
-    if plugins_list:
-        ArchiveResult.objects.filter(snapshot_id__in=existing_snapshot_ids, plugin__in=plugins_list).update(**reset_fields)
-    elif requested_plugins_by_snapshot:
-        snapshot_ids_by_plugin: dict[str, set[str]] = defaultdict(set)
-        for snapshot_id, plugin_names in requested_plugins_by_snapshot.items():
-            if snapshot_id in existing_snapshot_ids:
-                for plugin_name in plugin_names:
-                    snapshot_ids_by_plugin[plugin_name].add(snapshot_id)
-        for plugin_name, plugin_snapshot_ids in snapshot_ids_by_plugin.items():
-            ArchiveResult.objects.filter(snapshot_id__in=plugin_snapshot_ids, plugin=plugin_name).update(**reset_fields)
-    existing_rows = set(
-        ArchiveResult.objects.filter(
-            snapshot_id__in=existing_snapshot_ids,
-            plugin__in={plugin_name for _snapshot_id, plugin_name, _hook_name in requested_rows},
-        ).values_list("snapshot_id", "plugin", "hook_name"),
+    if rows_to_queue and plugins_list:
+        rows_to_reset_by_hook: dict[tuple[str, str], set[str]] = defaultdict(set)
+        for snapshot_id, plugin_name, hook_name in rows_to_queue:
+            rows_to_reset_by_hook[(plugin_name, hook_name)].add(snapshot_id)
+        for (plugin_name, hook_name), plugin_snapshot_ids in rows_to_reset_by_hook.items():
+            ArchiveResult.objects.filter(snapshot_id__in=plugin_snapshot_ids, plugin=plugin_name, hook_name=hook_name).update(
+                **reset_fields,
+            )
+    elif rows_to_queue and requested_plugins_by_snapshot:
+        snapshot_ids_by_hook: dict[tuple[str, str], set[str]] = defaultdict(set)
+        for snapshot_id, plugin_name, hook_name in rows_to_queue:
+            snapshot_ids_by_hook[(plugin_name, hook_name)].add(snapshot_id)
+        for (plugin_name, hook_name), plugin_snapshot_ids in snapshot_ids_by_hook.items():
+            ArchiveResult.objects.filter(snapshot_id__in=plugin_snapshot_ids, plugin=plugin_name, hook_name=hook_name).update(
+                **reset_fields,
+            )
+    existing_rows = (
+        {
+            (str(snapshot_id), plugin_name, hook_name)
+            for snapshot_id, plugin_name, hook_name in ArchiveResult.objects.filter(
+                snapshot_id__in=existing_snapshot_ids,
+                plugin__in={plugin_name for _snapshot_id, plugin_name, _hook_name in rows_to_queue},
+            ).values_list("snapshot_id", "plugin", "hook_name")
+        }
+        if rows_to_queue
+        else set()
     )
-    missing_rows = requested_rows - {(str(snapshot_id), plugin_name, hook_name) for snapshot_id, plugin_name, hook_name in existing_rows}
+    missing_rows = rows_to_queue - existing_rows
     if missing_rows:
         ArchiveResult.objects.bulk_create(
             [
@@ -271,7 +295,23 @@ def run_plugins(
             # also keep status=paused here: `retry_at` only asks the orchestrator
             # to process the queued plugin rows, and run_due_snapshot restores
             # retry_at=MAX afterward instead of resuming the snapshot lifecycle.
-            for snapshot in Snapshot.objects.filter(id__in=existing_snapshot_ids).only("id", "status", "modified_at"):
+            affected_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name, _hook_name in rows_to_queue}
+            if preserve_queued and queued_rows:
+                queued_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name, _hook_name in queued_rows}
+                affected_snapshot_ids.update(
+                    str(snapshot_id)
+                    for snapshot_id in Snapshot.objects.filter(id__in=queued_snapshot_ids)
+                    .filter(retry_at__gt=queue_at)
+                    .values_list("id", flat=True)
+                )
+                affected_snapshot_ids.update(
+                    str(snapshot_id)
+                    for snapshot_id in Snapshot.objects.filter(id__in=queued_snapshot_ids, retry_at__isnull=True).values_list(
+                        "id",
+                        flat=True,
+                    )
+                )
+            for snapshot in Snapshot.objects.filter(id__in=affected_snapshot_ids).only("id", "status", "modified_at"):
                 # Guard the read-time status so we never bump retry_at on a
                 # row that's been re-queued / started by a concurrent runner.
                 snapshot.safe_update(

archivebox/cli/archivebox_update.py

@@ -105,9 +105,17 @@ def reindex_snapshots(
     wait_for_turn=None,
 ) -> dict[str, Any]:
     from archivebox.cli.archivebox_extract import run_plugins
+    from archivebox.core.models import ArchiveResult
+    from abx_dl.models import discover_plugins
 
-    stats: dict[str, Any] = {"processed": 0, "queued": 0, "reindexed": 0, "snapshot_ids": []}
+    stats: dict[str, Any] = {"processed": 0, "requested": 0, "queued": 0, "skipped_queued": 0, "reindexed": 0, "snapshot_ids": []}
     records: list[dict[str, str]] = []
+    plugins_by_name = discover_plugins()
+    required_hooks_by_plugin = {
+        plugin_name: frozenset(hook.name for hook in plugins_by_name[plugin_name].filter_hooks("Snapshot"))
+        for plugin_name in search_plugins
+        if plugin_name in plugins_by_name
+    }
 
     total = snapshots.count()
     print(f"[*] Reindexing {total} snapshots with search plugins: {', '.join(search_plugins)}")
@@ -118,6 +126,31 @@ def reindex_snapshots(
         if wait_for_turn:
             wait_for_turn()
         batch_records = list(records)
+        snapshot_ids = {record["snapshot_id"] for record in batch_records}
+        plugin_names = {record["plugin"] for record in batch_records}
+        queued_rows = {
+            (str(snapshot_id), plugin_name, hook_name)
+            for snapshot_id, plugin_name, hook_name in ArchiveResult.objects.filter(
+                snapshot_id__in=snapshot_ids,
+                plugin__in=plugin_names,
+                status=ArchiveResult.StatusChoices.QUEUED,
+            ).values_list("snapshot_id", "plugin", "hook_name")
+        }
+        records_to_queue = []
+        for record in batch_records:
+            snapshot_id = record["snapshot_id"]
+            plugin_name = record["plugin"]
+            required_hooks = required_hooks_by_plugin.get(plugin_name, frozenset())
+            if required_hooks and all((snapshot_id, plugin_name, hook_name) in queued_rows for hook_name in required_hooks):
+                stats["skipped_queued"] += 1
+                continue
+            records_to_queue.append(record)
+        if not records_to_queue:
+            print(
+                f"    [{stats['processed']}/{total}] Already queued {len(batch_records)} index jobs",
+            )
+            records.clear()
+            return
         # `archivebox update --index-only` intentionally breaks the usual
         # "runner discovers work" rule by inserting synthetic queued
         # ArchiveResult rows for search backends. run_plugins() keeps this as
@@ -127,15 +160,17 @@ def reindex_snapshots(
         # finish.
         exit_code = run_plugins(
             args=(),
-            records=batch_records,
+            records=records_to_queue,
             wait=False,
             emit_results=False,
             show_progress=False,
+            preserve_queued=True,
         )
         if exit_code != 0:
             raise SystemExit(exit_code)
+        stats["queued"] += len(records_to_queue)
         print(
-            f"    [{stats['processed']}/{total}] Queued {len(batch_records)} index jobs for orchestrator",
+            f"    [{stats['processed']}/{total}] Queued {len(records_to_queue)} index jobs for orchestrator",
         )
         records.clear()
 
@@ -156,7 +191,7 @@ def reindex_snapshots(
                         "plugin": plugin_name,
                     },
                 )
-                stats["queued"] += 1
+                stats["requested"] += 1
             if len(records) >= batch_size:
                 run_batch()
         except KeyboardInterrupt as err:
@@ -366,15 +401,55 @@ def update(
                             after=after,
                             resume=resume,
                         )
-                        stats = reindex_snapshots(
-                            snapshots,
-                            search_plugins=search_plugins,
-                            batch_size=batch_size,
-                            collect_ids=is_filtered_update,
-                            wait_for_turn=wait_for_turn,
+                        from django.db.models import Exists, OuterRef, Q
+                        from django.utils import timezone
+                        from archivebox.core.models import ArchiveResult, Snapshot
+
+                        scoped_snapshot_ids = snapshots.order_by().values("id") if is_filtered_update else None
+                        queued_index_results = ArchiveResult.objects.filter(
+                            status=ArchiveResult.StatusChoices.QUEUED,
+                            plugin__in=search_plugins,
                         )
-                        print_index_stats(stats)
-                        touched_snapshot_ids.update(stats.get("snapshot_ids", []))
+                        if scoped_snapshot_ids is not None:
+                            queued_index_results = queued_index_results.filter(snapshot_id__in=scoped_snapshot_ids)
+
+                        if queued_index_results.exists():
+                            now = timezone.now()
+                            queued_result_for_snapshot = queued_index_results.filter(snapshot_id=OuterRef("pk"))
+                            snapshots_to_wake = (
+                                Snapshot.objects.filter(
+                                    status__in=(Snapshot.StatusChoices.SEALED, Snapshot.StatusChoices.PAUSED),
+                                )
+                                .annotate(
+                                    has_queued_index_result=Exists(queued_result_for_snapshot),
+                                )
+                                .filter(
+                                    has_queued_index_result=True,
+                                )
+                                .filter(
+                                    Q(retry_at__isnull=True) | Q(retry_at__gt=now),
+                                )
+                            )
+                            if scoped_snapshot_ids is not None:
+                                snapshots_to_wake = snapshots_to_wake.filter(id__in=scoped_snapshot_ids)
+                            woken_count = snapshots_to_wake.update(
+                                retry_at=now,
+                                modified_at=now,
+                            )
+                            print(
+                                "[*] Existing queued search index jobs found; "
+                                f"skipping backfill scan and waking {woken_count} snapshot(s) for the runner.",
+                            )
+                        else:
+                            stats = reindex_snapshots(
+                                snapshots,
+                                search_plugins=search_plugins,
+                                batch_size=batch_size,
+                                collect_ids=is_filtered_update,
+                                wait_for_turn=wait_for_turn,
+                            )
+                            print_index_stats(stats)
+                            touched_snapshot_ids.update(stats.get("snapshot_ids", []))
 
                 if do_run_until_idle and (do_index or not ran_post_migrate_runner):
                     # Search/index backfill intentionally queues targeted
@@ -904,7 +979,9 @@ def print_index_stats(stats: dict[str, Any]) -> None:
     print(f"""
 [green]Search Reindex Complete[/green]
   Scanned rows:      {stats["processed"]}
+  Requested jobs:    {stats.get("requested", stats["queued"])}
   Queued index jobs: {stats["queued"]}
+  Already queued:    {stats.get("skipped_queued", 0)}
 """)
 
 

archivebox/machine/models.py

@@ -548,6 +548,7 @@ class Binary(ModelWithHealthStats, ModelWithStateMachine):
 
     state_machine_name: str | None = "archivebox.machine.models.BinaryMachine"
     active_state: str = StatusChoices.QUEUED
+    warn_on_save_outside_runner = False
 
     objects = BinaryManager()  # pyright: ignore[reportIncompatibleVariableOverride]
 

archivebox/workers/models.py

@@ -65,6 +65,7 @@ class BaseModelWithStateMachine(models.Model):
     state_field_name: str
     state_machine_attr: str = "sm"
     bind_events_as_methods: bool = False
+    warn_on_save_outside_runner: ClassVar[bool] = True
 
     active_state: ObjectState
     retry_at_field_name: str
@@ -297,7 +298,7 @@ class BaseModelWithStateMachine(models.Model):
         from archivebox.machine.models import Process
 
         process = Process.current()
-        if process.process_type != Process.TypeChoices.ORCHESTRATOR:
+        if self.warn_on_save_outside_runner and process.process_type != Process.TypeChoices.ORCHESTRATOR:
             root_type = getattr(process.root, "process_type", None)
             if root_type != Process.TypeChoices.ORCHESTRATOR:
                 caller = "<unknown>"