ArchiveBox/archivebox/services/crawl_service.py
2026-06-01 00:08:27 -07:00

106 lines
3.8 KiB
Python

from __future__ import annotations
from datetime import timedelta
from django.utils import timezone
from abx_dl.events import CrawlCleanupEvent, CrawlCompletedEvent, CrawlSetupEvent, CrawlStartEvent
from abx_dl.services.base import BaseService
from archivebox.workers.models import ACTIVE_STATE_LEASE_SECONDS
class CrawlService(BaseService):
LISTENS_TO = [CrawlSetupEvent, CrawlStartEvent, CrawlCleanupEvent, CrawlCompletedEvent]
EMITS = []
def __init__(self, bus, *, crawl_id: str):
self.crawl_id = crawl_id
super().__init__(bus)
self.bus.on(CrawlSetupEvent, self.on_CrawlSetupEvent__save_to_db)
self.bus.on(CrawlStartEvent, self.on_CrawlStartEvent__save_to_db)
self.bus.on(CrawlCleanupEvent, self.on_CrawlCleanupEvent__save_to_db)
self.bus.on(CrawlCompletedEvent, self.on_CrawlCompletedEvent__save_to_db)
async def on_CrawlSetupEvent__save_to_db(self, event: CrawlSetupEvent) -> None:
from archivebox.crawls.models import Crawl
await (
Crawl.objects.filter(id=self.crawl_id)
.exclude(
status__in=Crawl.INACTIVE_STATES,
)
.aupdate(
status=Crawl.StatusChoices.STARTED,
retry_at=timezone.now() + timedelta(seconds=ACTIVE_STATE_LEASE_SECONDS),
modified_at=timezone.now(),
)
)
async def on_CrawlStartEvent__save_to_db(self, event: CrawlStartEvent) -> None:
from archivebox.crawls.models import Crawl
await (
Crawl.objects.filter(id=self.crawl_id)
.exclude(
status__in=Crawl.INACTIVE_STATES,
)
.aupdate(
status=Crawl.StatusChoices.STARTED,
retry_at=timezone.now() + timedelta(seconds=ACTIVE_STATE_LEASE_SECONDS),
modified_at=timezone.now(),
)
)
async def on_CrawlCleanupEvent__save_to_db(self, event: CrawlCleanupEvent) -> None:
from archivebox.crawls.models import Crawl
# Cleanup is still inside the active crawl lifecycle. Snapshot hooks may
# have just written discovery output that the runner consumes before the
# completion phase, so only CrawlCompleted/finalize_run_state makes the
# final sealed-vs-requeue decision.
await (
Crawl.objects.filter(id=self.crawl_id)
.exclude(
status__in=Crawl.INACTIVE_STATES,
)
.aupdate(
status=Crawl.StatusChoices.STARTED,
retry_at=timezone.now(),
modified_at=timezone.now(),
)
)
async def on_CrawlCompletedEvent__save_to_db(self, event: CrawlCompletedEvent) -> None:
from archivebox.crawls.models import Crawl
from archivebox.core.models import Snapshot
crawl = await Crawl.objects.aget(id=self.crawl_id)
if crawl.is_paused or crawl.status == Crawl.StatusChoices.SEALED:
return
is_finished = not await crawl.snapshot_set.filter(status__in=Snapshot.OPEN_STATES).aexists()
if not is_finished:
await (
Crawl.objects.filter(id=self.crawl_id)
.exclude(
status__in=Crawl.INACTIVE_STATES,
)
.aupdate(
status=Crawl.StatusChoices.STARTED,
retry_at=timezone.now(),
modified_at=timezone.now(),
)
)
return
await (
Crawl.objects.filter(id=self.crawl_id)
.exclude(
status__in=Crawl.INACTIVE_STATES,
)
.aupdate(
status=Crawl.StatusChoices.SEALED,
retry_at=None,
modified_at=timezone.now(),
)
)