fix: harden concurrent snapshot retries

archivebox/api/v1_cli.py

@@ -165,7 +165,7 @@ def cli_add(request: HttpRequest, args: AddCommandSchema):
         "crawl_id": str(crawl.id),
         "num_snapshots": len(snapshot_ids),
         "snapshot_ids": snapshot_ids,
-        "queued_urls": args.urls,
+        "queued_urls": args.urls if isinstance(submitted_urls, str) else submitted_urls,
     }
     stdout = request.__dict__.get("stdout")
     stderr = request.__dict__.get("stderr")

archivebox/core/models.py

@@ -730,21 +730,32 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW
         plugin_names = sorted({name.strip() for name in plugins if name.strip()})
         if not plugin_names:
             return False
-        from django.db import transaction
-
         retry_at = when or timezone.now()
-        with transaction.atomic():
-            current = type(self).objects.select_for_update().select_related("crawl").get(pk=self.pk)
+        for _attempt in range(3):
+            current = type(self).objects.select_related("crawl").get(pk=self.pk)
             pending_plugins = {str(name).strip() for name in (current.config or {}).get("RETRY_PLUGINS", []) if str(name).strip()}
             config = {**(current.config or {}), "RETRY_PLUGINS": sorted(pending_plugins | set(plugin_names))}
             status = current.status if current.status == self.StatusChoices.SEALED else self.StatusChoices.QUEUED
-            type(self).objects.filter(pk=self.pk).update(
-                config=config,
-                status=status,
-                retry_at=retry_at,
-                modified_at=timezone.now(),
+            updated = (
+                type(self)
+                .objects.filter(
+                    pk=self.pk,
+                    config=current.config,
+                    status=current.status,
+                    retry_at=current.retry_at,
+                )
+                .update(
+                    config=config,
+                    status=status,
+                    retry_at=retry_at,
+                    modified_at=timezone.now(),
+                )
             )
-            crawl = current.crawl
+            if updated:
+                crawl = current.crawl
+                break
+        else:
+            return False
 
         self.refresh_from_db()
         if status in self.RUNNABLE_STATES and self.crawl_id:

archivebox/services/runner.py

@@ -1217,20 +1217,18 @@ def run_snapshot_maintenance(snapshot_id: str, *, output_dir: Path | None = None
         next_retry_at = timezone.now()
     else:
         next_retry_at = None
-    snapshot.retry_at = next_retry_at
     if snapshot.fs_migration_needed:
-        snapshot.save(update_fields=["retry_at", "modified_at"])
-    else:
-        updated = snapshot.safe_update(
-            {"retry_at": next_retry_at},
-            refresh=False,
-            extra_filter={
-                "status": snapshot.StatusChoices.SEALED,
-                "retry_at": current_retry_at,
-            },
-        )
-        if not updated:
-            return False
+        snapshot.migrate_filesystem_to_current_version()
+    updated = snapshot.safe_update(
+        {"retry_at": next_retry_at},
+        refresh=False,
+        extra_filter={
+            "status": snapshot.status,
+            "retry_at": current_retry_at,
+        },
+    )
+    if not updated:
+        return False
     snapshot.write_index_jsonl(output_dir=output_dir)
     return True
 
@@ -1370,13 +1368,14 @@ def _run_due_snapshot_locked(snapshot, *, lock_seconds: int, interactive_interru
         if not Snapshot.claim_for_worker(snapshot, lock_seconds=lock_seconds):
             return False
         snapshot.refresh_from_db()
+        owned_retry_at = snapshot.retry_at
         snapshot.finalize_completed_upload_results()
-        retry_plugins = [str(name).strip() for name in (snapshot.config or {}).get("RETRY_PLUGINS", []) if str(name).strip()]
         if snapshot.fs_migration_needed:
-            # Preserve the claimed retry_at lease while the idempotent
-            # migration runs so a pending plugin request remains discoverable.
-            snapshot.save(update_fields=["retry_at", "modified_at"])
-            snapshot.refresh_from_db()
+            snapshot.migrate_filesystem_to_current_version()
+        snapshot.refresh_from_db()
+        if snapshot.status != Snapshot.StatusChoices.SEALED or snapshot.retry_at != owned_retry_at:
+            return True
+        retry_plugins = [str(name).strip() for name in (snapshot.config or {}).get("RETRY_PLUGINS", []) if str(name).strip()]
         if retry_plugins:
             _runner_console_line(crawl_id=snapshot.crawl_id, snapshot=snapshot)
             run_crawl(
@@ -1402,7 +1401,7 @@ def _run_due_snapshot_locked(snapshot, *, lock_seconds: int, interactive_interru
         # Migrate before abx-dl writes new hook outputs. The claimed Snapshot
         # lease remains in place and the idempotent migration persists its
         # indexed fs_version marker only after copy/verification/cleanup.
-        snapshot.save(update_fields=["retry_at", "modified_at"])
+        snapshot.migrate_filesystem_to_current_version()
         snapshot.refresh_from_db()
     if snapshot.status == Snapshot.StatusChoices.QUEUED:
         snapshot.start_processing()

archivebox/services/snapshot_service.py

@@ -154,7 +154,12 @@ class SnapshotService(BaseService):
     async def on_SnapshotCompletedEvent(self, event: SnapshotCompletedEvent) -> None:
         ownership = self._run_ownership.pop(str(event.snapshot_id), None)
         if ownership is None:
-            return
+            from archivebox.core.models import Snapshot
+
+            snapshot = await Snapshot.objects.only("status", "retry_at").filter(id=event.snapshot_id, crawl_id=self.crawl_id).afirst()
+            if snapshot is None:
+                return
+            ownership = (snapshot.retry_at, snapshot.status == Snapshot.StatusChoices.SEALED, [])
         owned_retry_at, was_sealed, retry_plugins = ownership
         await sync_to_async(finalize_completed_snapshot, thread_sensitive=True)(
             event.snapshot_id,

archivebox/tests/test_api_v1_cli_add.py

@@ -285,6 +285,7 @@ def test_api_cli_add_filters_invalid_items_from_multi_url_batch(client, tmp_path
     )
 
     assert response.status_code == 200, response.content
+    assert response.json()["result"]["queued_urls"] == [submitted_url]
     assert Crawl.objects.get().urls == submitted_url