mirror of
https://github.com/ArchiveBox/ArchiveBox.git
synced 2026-09-12 19:50:57 +05:00
Unify plugin runtime and queue lifecycle boundaries
This commit is contained in:
parent
402c983f61
commit
3fc67705a5
@ -342,21 +342,10 @@ def _parse_archiveresult_upload_int(value: str, field_name: str, *, default: int
|
||||
|
||||
|
||||
def _summarize_archiveresult_output_files(output_files: dict[str, dict[str, Any]]) -> tuple[int, str]:
|
||||
mime_sizes: dict[str, int] = defaultdict(int)
|
||||
total_size = 0
|
||||
for metadata in output_files.values():
|
||||
if not isinstance(metadata, dict):
|
||||
continue
|
||||
try:
|
||||
size = max(int(metadata.get("size") or 0), 0)
|
||||
except (TypeError, ValueError):
|
||||
size = 0
|
||||
mime_type = str(metadata.get("mimetype") or "").strip()
|
||||
total_size += size
|
||||
if mime_type and size:
|
||||
mime_sizes[mime_type] += size
|
||||
output_mimetypes = ",".join(mime for mime, _size in sorted(mime_sizes.items(), key=lambda item: item[1], reverse=True))
|
||||
return total_size, output_mimetypes
|
||||
from abx_dl.output_files import OutputManifest
|
||||
|
||||
manifest = OutputManifest.from_value(output_files)
|
||||
return manifest.total_size, ",".join(manifest.mimetypes)
|
||||
|
||||
|
||||
def _get_snapshot_by_ref(snapshot_id: str):
|
||||
|
||||
@ -264,7 +264,7 @@ def run_plugins(
|
||||
if requested_pairs:
|
||||
# Search indexing on a sealed Snapshot is the only targeted plugin
|
||||
# allowed to bypass the normal lifecycle. Every other requested
|
||||
# plugin requeues its Snapshot through the unified state machine.
|
||||
# plugin requeues its Snapshot through the unified lifecycle.
|
||||
affected_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name in rows_to_queue}
|
||||
if preserve_queued and queued_rows:
|
||||
queued_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name in queued_rows}
|
||||
|
||||
@ -54,7 +54,7 @@ def _resolve_install_targets(
|
||||
|
||||
|
||||
def _install_raw_binary_names(binary_names: list[str], binproviders: str) -> None:
|
||||
"""Install user-requested standalone binaries through the Binary state machine."""
|
||||
"""Install user-requested standalone binaries through the Binary lifecycle."""
|
||||
from django.utils import timezone
|
||||
|
||||
from archivebox.machine.models import Binary, Machine, _canonical_binary_name
|
||||
@ -179,7 +179,7 @@ def install(binaries: tuple[str, ...] = (), binproviders: str = "*", dry_run: bo
|
||||
run_install(plugin_names=install_plugin_names or None)
|
||||
|
||||
if raw_binary_names:
|
||||
print(f"[+] Running direct binary installer via ArchiveBox binary state machine: {', '.join(raw_binary_names)}")
|
||||
print(f"[+] Running direct binary installer via ArchiveBox binary lifecycle: {', '.join(raw_binary_names)}")
|
||||
print()
|
||||
_install_raw_binary_names(raw_binary_names, binproviders)
|
||||
|
||||
|
||||
@ -48,7 +48,7 @@ def pluginmap(
|
||||
Show the current abx-dl event phases and their associated plugin hooks.
|
||||
|
||||
This command reflects the new bus-driven runtime, not the legacy ArchiveBox
|
||||
state-machine executor. Event names are normalized to hook prefixes by
|
||||
event runtime. Event names are normalized to hook prefixes by
|
||||
stripping a trailing `Event`, then ArchiveBox checks whether any matching
|
||||
`on_{EventFamily}__*` scripts actually exist.
|
||||
"""
|
||||
|
||||
@ -81,7 +81,7 @@ def reindex_snapshots(
|
||||
|
||||
# Search backfill is the one maintenance hook allowed to execute without
|
||||
# reopening a Snapshot. Restrict that exception to already-sealed rows;
|
||||
# every open lifecycle state remains owned by the normal state machine.
|
||||
# every open lifecycle state remains owned by the normal runner lifecycle.
|
||||
snapshots = snapshots.filter(status=Snapshot.StatusChoices.SEALED)
|
||||
|
||||
stats: dict[str, Any] = {"processed": 0, "requested": 0, "queued": 0, "skipped_queued": 0, "reindexed": 0, "snapshot_ids": []}
|
||||
|
||||
@ -16,7 +16,7 @@ from pathlib import Path
|
||||
from typing import Any, ClassVar, cast
|
||||
from urllib.parse import quote, urlparse
|
||||
|
||||
from abx_plugins.plugins.base.utils import BASE_CONFIG_PATH, build_config_model, resolve_plugin_configs
|
||||
from abx_plugins.plugins.base.utils import build_config_model
|
||||
from django.db import DatabaseError
|
||||
from pydantic import BaseModel, Field, PrivateAttr, create_model, field_validator, model_validator
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
@ -601,13 +601,9 @@ def _explicit_plugin_enabled_keys(config: Mapping[str, object]) -> set[str]:
|
||||
|
||||
|
||||
def _discover_plugin_config_schemas() -> PluginSchemaDocuments:
|
||||
from archivebox.plugins.discovery import discover_plugin_configs
|
||||
from archivebox.plugins.discovery import get_plugin_config_resolver
|
||||
|
||||
schemas: PluginSchemaDocuments = {}
|
||||
if BASE_CONFIG_PATH.exists():
|
||||
schemas["base"] = json.loads(BASE_CONFIG_PATH.read_text())
|
||||
schemas.update(discover_plugin_configs())
|
||||
return schemas
|
||||
return get_plugin_config_resolver().schemas
|
||||
|
||||
|
||||
def _plugin_config_properties(plugin_schemas: PluginSchemaDocuments) -> dict[str, dict[str, Any]]:
|
||||
@ -1195,7 +1191,6 @@ def get_config(
|
||||
|
||||
explicit_plugin_enabled_keys: set[str] = set()
|
||||
if resolve_plugins:
|
||||
plugin_schemas = {plugin_name: schema for plugin_name, schema in PLUGIN_CONFIG_SCHEMAS.items() if isinstance(schema, dict)}
|
||||
plugin_global_config = {key: str(value) if isinstance(value, Path) else value for key, value in config_data.items()}
|
||||
crawl_selected_plugins = crawl_config_base and bool(_normalize_plugins_config_value(dict(crawl.config or {}).get("PLUGINS")))
|
||||
# A frozen crawl-level PLUGINS selector is the exact extractor set for
|
||||
@ -1219,8 +1214,9 @@ def get_config(
|
||||
**_plugin_user_config(_plugin_input_config(file_config)),
|
||||
**plugin_user_config,
|
||||
}
|
||||
plugin_sections = resolve_plugin_configs(
|
||||
plugin_schemas,
|
||||
from archivebox.plugins.discovery import get_plugin_config_resolver
|
||||
|
||||
plugin_sections = get_plugin_config_resolver().resolve(
|
||||
global_config=plugin_global_config,
|
||||
user_config=plugin_user_config,
|
||||
environ={},
|
||||
|
||||
@ -10,7 +10,6 @@ class CoreConfig(AppConfig):
|
||||
|
||||
def ready(self):
|
||||
"""Register the archivebox.core.admin_site as the main django admin site"""
|
||||
import sys
|
||||
from django.utils.autoreload import DJANGO_AUTORELOAD_ENV
|
||||
|
||||
from archivebox.core.admin_site import register_admin_site
|
||||
@ -28,11 +27,6 @@ class CoreConfig(AppConfig):
|
||||
|
||||
pre_save.connect(truncate_overlong_charfields, dispatch_uid="archivebox_truncate_overlong_charfields")
|
||||
|
||||
# Import models to register state machines with the registry
|
||||
# Skip during makemigrations to avoid premature state machine access
|
||||
if "makemigrations" not in sys.argv:
|
||||
from archivebox.core import models # noqa: F401
|
||||
|
||||
def _should_prepare_runtime() -> bool:
|
||||
if os.environ.get("ARCHIVEBOX_RUNSERVER") == "1":
|
||||
if os.environ.get("ARCHIVEBOX_AUTORELOAD") == "1":
|
||||
|
||||
@ -12,9 +12,9 @@ from archivebox.core.widgets import TagEditorWidget, URLFiltersWidget
|
||||
from archivebox.crawls.schedule_util import validate_schedule
|
||||
from archivebox.misc.util import URL_REGEX, find_all_urls, parse_filesize_to_bytes
|
||||
from archivebox.personas.models import Persona
|
||||
from archivebox.plugins.discovery import get_plugins
|
||||
from archivebox.plugins.discovery import get_plugin_catalog
|
||||
from archivebox.plugins.forms import (
|
||||
PLUGIN_GROUP_DEFINITIONS,
|
||||
PLUGIN_GROUPS,
|
||||
TIMEOUT_INPUT_PATTERN,
|
||||
PluginConfigFormMixin,
|
||||
get_choice_field,
|
||||
@ -333,10 +333,9 @@ class AddLinkForm(PluginConfigFormMixin, forms.Form):
|
||||
if self.can_override_crawl_config:
|
||||
self.build_plugin_groups(get_config(persona=selected_persona) if selected_persona else get_config())
|
||||
else:
|
||||
all_plugins = get_plugins()
|
||||
for field_name, *_rest, plugin_names in PLUGIN_GROUP_DEFINITIONS:
|
||||
get_choice_field(self, field_name).choices = [(p, p) for p in all_plugins if p in plugin_names]
|
||||
get_choice_field(self, "other_plugins").choices = [(p, p) for p in all_plugins]
|
||||
grouped_plugins = get_plugin_catalog().groups()
|
||||
for category, field_name, _title in PLUGIN_GROUPS:
|
||||
get_choice_field(self, field_name).choices = [(plugin.name, plugin.name) for plugin in grouped_plugins.get(category, [])]
|
||||
self.plugin_groups = []
|
||||
|
||||
def clean(self):
|
||||
|
||||
@ -21,7 +21,6 @@ from django.utils import timezone
|
||||
from django.utils.functional import cached_property
|
||||
from django.utils.safestring import mark_safe
|
||||
from django.utils.text import slugify
|
||||
from statemachine import State, registry
|
||||
|
||||
from archivebox.base_models.models import (
|
||||
ModelWithConfig,
|
||||
@ -54,7 +53,7 @@ from archivebox.plugins.discovery import (
|
||||
get_plugins,
|
||||
)
|
||||
from archivebox.uuid_compat import CompactUUIDField, uuid7
|
||||
from archivebox.workers.models import ACTIVE_STATE_LEASE_SECONDS, RETRY_AT_MAX, BaseStateMachine, ModelWithStateMachine
|
||||
from archivebox.workers.models import ACTIVE_STATE_LEASE_SECONDS, RETRY_AT_MAX, ModelWithQueue
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from archivebox.config.common import ArchiveBoxBaseConfig
|
||||
@ -523,7 +522,7 @@ class SnapshotManager(models.Manager.from_queryset(SnapshotQuerySet)): # ty: ig
|
||||
return self.get_queryset().delete()
|
||||
|
||||
|
||||
class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWithNotes, ModelWithHealthStats, ModelWithStateMachine):
|
||||
class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWithNotes, ModelWithHealthStats, ModelWithQueue):
|
||||
BROWSER_EXTENSION_UPLOAD_HOOK_NAME = "on_Snapshot__archivebox_browser_extension_upload"
|
||||
|
||||
INTERNAL_INPUT_URL = "archivebox://internal"
|
||||
@ -563,10 +562,10 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW
|
||||
help_text="Current hook step being executed (0-9). Used for sequential hook execution.",
|
||||
)
|
||||
|
||||
retry_at = ModelWithStateMachine.RetryAtField(default=timezone.now)
|
||||
status = ModelWithStateMachine.StatusField(
|
||||
choices=ModelWithStateMachine.StatusChoices,
|
||||
default=ModelWithStateMachine.StatusChoices.QUEUED,
|
||||
retry_at = ModelWithQueue.RetryAtField(default=timezone.now)
|
||||
status = ModelWithQueue.StatusField(
|
||||
choices=ModelWithQueue.StatusChoices,
|
||||
default=ModelWithQueue.StatusChoices.QUEUED,
|
||||
)
|
||||
config = models.JSONField(default=dict, null=False, blank=False, editable=True)
|
||||
permissions = models.GeneratedField(
|
||||
@ -587,10 +586,13 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW
|
||||
|
||||
tags = models.ManyToManyField(Tag, blank=True, through=SnapshotTag, related_name="snapshot_set", through_fields=("snapshot", "tag"))
|
||||
|
||||
state_machine_name = "archivebox.core.models.SnapshotMachine"
|
||||
state_field_name = "status"
|
||||
retry_at_field_name = "retry_at"
|
||||
StatusChoices = ModelWithStateMachine.StatusChoices
|
||||
StatusChoices = ModelWithQueue.StatusChoices
|
||||
INITIAL_STATE = StatusChoices.QUEUED
|
||||
ACTIVE_STATE = StatusChoices.STARTED
|
||||
FINAL_STATES = (StatusChoices.SEALED,)
|
||||
FINAL_OR_ACTIVE_STATES = (*FINAL_STATES, ACTIVE_STATE)
|
||||
active_state = StatusChoices.STARTED
|
||||
delete_after_final_statuses = (StatusChoices.SEALED,)
|
||||
RUNNABLE_STATES = (StatusChoices.QUEUED, StatusChoices.STARTED)
|
||||
@ -603,11 +605,6 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW
|
||||
objects = SnapshotManager()
|
||||
archiveresult_set: models.Manager["ArchiveResult"]
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@property
|
||||
def sm(self) -> "SnapshotMachine": ...
|
||||
|
||||
def add_tag_ids(self, tag_ids: Iterable[int | str]) -> None:
|
||||
tag_ids = [tag_id for tag_id in dict.fromkeys(tag_ids) if tag_id]
|
||||
if not tag_ids:
|
||||
@ -623,7 +620,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW
|
||||
ModelWithConfig.Meta,
|
||||
ModelWithNotes.Meta,
|
||||
ModelWithHealthStats.Meta,
|
||||
ModelWithStateMachine.Meta,
|
||||
ModelWithQueue.Meta,
|
||||
):
|
||||
app_label = "core"
|
||||
verbose_name = "Snapshot"
|
||||
@ -744,7 +741,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW
|
||||
Crawl.pause()/cancel() only wake child rows. The runner claims each due
|
||||
Snapshot and lets this method perform the actual child transition, so
|
||||
cancellation stays fast and Snapshot cleanup still runs from the normal
|
||||
state-machine owner.
|
||||
lifecycle owner.
|
||||
"""
|
||||
parent_status = Crawl.objects.filter(id=self.crawl_id).values_list("status", flat=True).first()
|
||||
if parent_status == Crawl.StatusChoices.SEALED and self.status != self.StatusChoices.SEALED:
|
||||
@ -753,7 +750,7 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW
|
||||
self.refresh_from_db()
|
||||
parent_status = Crawl.objects.filter(id=self.crawl_id).values_list("status", flat=True).first()
|
||||
if parent_status == Crawl.StatusChoices.SEALED and self.status != self.StatusChoices.SEALED:
|
||||
self.sm.seal()
|
||||
self.seal()
|
||||
return True
|
||||
|
||||
if parent_status == Crawl.StatusChoices.PAUSED and self.status not in (self.StatusChoices.PAUSED, self.StatusChoices.SEALED):
|
||||
@ -809,9 +806,66 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW
|
||||
reset_count += 1
|
||||
return reset_count, running_count
|
||||
|
||||
def start_processing(self) -> bool:
|
||||
"""Atomically move a claimed queued Snapshot into its active lease."""
|
||||
owned_retry_at = self.retry_at
|
||||
now = timezone.now()
|
||||
lease_until = now + timedelta(seconds=ACTIVE_STATE_LEASE_SECONDS)
|
||||
updated = (
|
||||
type(self)
|
||||
.objects.filter(
|
||||
pk=self.pk,
|
||||
retry_at=owned_retry_at,
|
||||
status=self.StatusChoices.QUEUED,
|
||||
)
|
||||
.update(
|
||||
status=self.StatusChoices.STARTED,
|
||||
retry_at=lease_until,
|
||||
modified_at=now,
|
||||
)
|
||||
)
|
||||
self.refresh_from_db()
|
||||
return updated == 1
|
||||
|
||||
def seal(self) -> bool:
|
||||
"""Atomically finalize this Snapshot and reconcile its output metadata."""
|
||||
if self.status == self.StatusChoices.SEALED:
|
||||
return True
|
||||
now = timezone.now()
|
||||
updated = (
|
||||
type(self)
|
||||
.objects.filter(
|
||||
pk=self.pk,
|
||||
retry_at=self.retry_at,
|
||||
status__in=self.OPEN_STATES,
|
||||
)
|
||||
.update(
|
||||
status=self.StatusChoices.SEALED,
|
||||
retry_at=None,
|
||||
modified_at=now,
|
||||
)
|
||||
)
|
||||
self.refresh_from_db()
|
||||
if updated == 1:
|
||||
self.finalize_output_metadata()
|
||||
return updated == 1
|
||||
|
||||
def advance_lifecycle(self) -> bool:
|
||||
"""Advance one explicit lifecycle step after the runner claims this row."""
|
||||
if self.status == self.StatusChoices.PAUSED:
|
||||
return False
|
||||
if self.status == self.StatusChoices.QUEUED:
|
||||
results = self.archiveresult_set.all()
|
||||
if results.exists() and not results.exclude(status__in=ArchiveResult.FINAL_STATES).exists():
|
||||
return self.seal()
|
||||
return bool(self.url) and self.start_processing()
|
||||
if self.status == self.StatusChoices.STARTED and self.is_finished_processing():
|
||||
return self.seal()
|
||||
return False
|
||||
|
||||
def cancel(self) -> None:
|
||||
if self.status != self.StatusChoices.SEALED:
|
||||
self.sm.seal()
|
||||
self.seal()
|
||||
|
||||
def get_delete_after_config_value(self):
|
||||
from archivebox.config.common import resolve_delete_after_config_value
|
||||
@ -2728,11 +2782,11 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW
|
||||
"""
|
||||
return self.create_pending_archiveresults()
|
||||
|
||||
def cleanup(self):
|
||||
def finalize_output_metadata(self) -> None:
|
||||
"""
|
||||
Clean up background ArchiveResult hooks and empty results.
|
||||
|
||||
Called by the state machine when entering the 'sealed' state.
|
||||
Called after entering the sealed state.
|
||||
Reconcile late background outputs and hydrate result metadata.
|
||||
"""
|
||||
# Clean up .pid files from output directory.
|
||||
@ -2742,11 +2796,10 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW
|
||||
for pid_file in output_dir.glob("**/*.pid"):
|
||||
pid_file.unlink(missing_ok=True)
|
||||
|
||||
# Update all background ArchiveResults from filesystem in case
|
||||
# output arrived late. If there is no snapshot directory, there is
|
||||
# no filesystem output to reconcile and no reason to hit this query.
|
||||
# Reconcile late background output without re-running hook-record
|
||||
# dispatch. The abx-dl event projector is the sole status owner.
|
||||
for ar in self.archiveresult_set.filter(hook_name__contains=".bg."):
|
||||
ar.update_from_output()
|
||||
ar.update_output_metadata_from_filesystem(snapshot_dir=output_dir)
|
||||
else:
|
||||
return
|
||||
|
||||
@ -3760,162 +3813,6 @@ class Snapshot(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelW
|
||||
return dt.strftime("%Y-%m-%d %H:%M:%S") if dt else None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Snapshot State Machine
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class SnapshotMachine(BaseStateMachine):
|
||||
"""
|
||||
State machine for managing Snapshot lifecycle.
|
||||
|
||||
Hook Lifecycle:
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ QUEUED State │
|
||||
│ • Waiting for snapshot to be ready │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓ tick() when can_start()
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ STARTED State → enter_started() │
|
||||
│ 1. snapshot.run() │
|
||||
│ • discover_hooks('Snapshot') → finds all plugin hooks │
|
||||
│ • create_pending_archiveresults() → creates ONE │
|
||||
│ ArchiveResult per hook (NO execution yet) │
|
||||
│ 2. The shared abx-dl runner executes hooks and the │
|
||||
│ projector updates ArchiveResult rows from events │
|
||||
│ 3. Advance through steps 0-9 as foreground hooks complete │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓ tick() when is_finished()
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ SEALED State → enter_sealed() │
|
||||
│ • cleanup() → kills any background hooks still running │
|
||||
│ • Set retry_at=None (no more processing) │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
|
||||
https://github.com/ArchiveBox/ArchiveBox/wiki/ArchiveBox-Architecture-Diagrams
|
||||
"""
|
||||
|
||||
model_attr_name = "snapshot"
|
||||
|
||||
# States
|
||||
queued = State(value=Snapshot.StatusChoices.QUEUED, initial=True)
|
||||
started = State(value=Snapshot.StatusChoices.STARTED)
|
||||
paused = State(value=Snapshot.StatusChoices.PAUSED)
|
||||
sealed = State(value=Snapshot.StatusChoices.SEALED, final=True)
|
||||
|
||||
# Tick Event (polled by workers)
|
||||
tick = (
|
||||
queued.to(sealed, cond="has_finished_archive_results")
|
||||
| queued.to.itself(unless="can_start")
|
||||
| queued.to(started, cond="can_start")
|
||||
| started.to(sealed, cond="is_finished")
|
||||
| paused.to.itself()
|
||||
)
|
||||
|
||||
# Manual event (can also be triggered by last ArchiveResult finishing)
|
||||
seal = queued.to(sealed) | started.to(sealed) | paused.to(sealed)
|
||||
pause_requested = queued.to(paused) | started.to(paused)
|
||||
resume_requested = paused.to(queued)
|
||||
|
||||
snapshot: Snapshot
|
||||
|
||||
def can_start(self) -> bool:
|
||||
can_start = bool(self.snapshot.url)
|
||||
return can_start
|
||||
|
||||
def is_finished(self) -> bool:
|
||||
"""Check if all ArchiveResults for this snapshot are finished."""
|
||||
return self.snapshot.is_finished_processing()
|
||||
|
||||
def has_finished_archive_results(self) -> bool:
|
||||
"""A queued snapshot with only final projected rows was interrupted after hook completion."""
|
||||
results = self.snapshot.archiveresult_set.all()
|
||||
return results.exists() and not results.exclude(status__in=ArchiveResult.FINAL_STATES).exists()
|
||||
|
||||
@queued.enter
|
||||
def enter_queued(self):
|
||||
self.snapshot.update_and_requeue(
|
||||
retry_at=timezone.now(),
|
||||
status=Snapshot.StatusChoices.QUEUED,
|
||||
)
|
||||
|
||||
@paused.enter
|
||||
def enter_paused(self):
|
||||
self.snapshot.safe_update(
|
||||
{
|
||||
"retry_at": RETRY_AT_MAX,
|
||||
"status": Snapshot.StatusChoices.PAUSED,
|
||||
},
|
||||
extra_filter={"status__in": Snapshot.RUNNABLE_STATES},
|
||||
)
|
||||
|
||||
@started.enter
|
||||
def enter_started(self):
|
||||
"""Just mark as started. The shared runner creates ArchiveResults and runs hooks."""
|
||||
owned_retry_at = self.snapshot.retry_at
|
||||
now = timezone.now()
|
||||
lease_until = now + timedelta(seconds=ACTIVE_STATE_LEASE_SECONDS)
|
||||
# The runner owns queued Snapshot startup through retry_at. Creating
|
||||
# pending ArchiveResult rows immediately before tick() can touch
|
||||
# Snapshot.modified_at, so using modified_at CAS here would reject the
|
||||
# legitimate owner. Keep the write to the scheduler columns only.
|
||||
updated = Snapshot.objects.filter(
|
||||
pk=self.snapshot.pk,
|
||||
retry_at=owned_retry_at,
|
||||
status=Snapshot.StatusChoices.QUEUED,
|
||||
).update(
|
||||
status=Snapshot.StatusChoices.STARTED,
|
||||
retry_at=lease_until,
|
||||
modified_at=now,
|
||||
)
|
||||
if updated != 1:
|
||||
self.snapshot.refresh_from_db()
|
||||
return
|
||||
self.snapshot.status = Snapshot.StatusChoices.STARTED
|
||||
self.snapshot.retry_at = lease_until
|
||||
self.snapshot.modified_at = now
|
||||
|
||||
@sealed.enter
|
||||
def enter_sealed(self):
|
||||
now = timezone.now()
|
||||
owned_retry_at = self.snapshot.retry_at
|
||||
# The runner owns this row via retry_at. Commit the final lifecycle
|
||||
# state before cleanup so late projectors can update metadata without
|
||||
# tripping a modified_at CAS while the row still looks QUEUED/STARTED.
|
||||
updated = (
|
||||
type(self.snapshot)
|
||||
.objects.filter(
|
||||
pk=self.snapshot.pk,
|
||||
retry_at=owned_retry_at,
|
||||
status__in=[
|
||||
Snapshot.StatusChoices.QUEUED,
|
||||
Snapshot.StatusChoices.STARTED,
|
||||
Snapshot.StatusChoices.PAUSED,
|
||||
],
|
||||
)
|
||||
.update(
|
||||
status=Snapshot.StatusChoices.SEALED,
|
||||
retry_at=None,
|
||||
modified_at=now,
|
||||
)
|
||||
)
|
||||
if updated != 1:
|
||||
self.snapshot.refresh_from_db()
|
||||
return
|
||||
|
||||
self.snapshot.status = Snapshot.StatusChoices.SEALED
|
||||
self.snapshot.retry_at = None
|
||||
self.snapshot.modified_at = now
|
||||
|
||||
# Clean up background hooks after the final state is visible in DB.
|
||||
self.snapshot.cleanup()
|
||||
|
||||
# Crawl finalization is handled by the runner/CrawlService cleanup
|
||||
# phase. Sealing the parent crawl here races recursive discovery:
|
||||
# Snapshot hooks can write urls.jsonl just before this state transition,
|
||||
# and the runner still needs to enqueue those child snapshots.
|
||||
|
||||
|
||||
class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes):
|
||||
class StatusChoices(models.TextChoices):
|
||||
QUEUED = "queued", "Queued"
|
||||
@ -4387,51 +4284,9 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes):
|
||||
|
||||
@staticmethod
|
||||
def _normalize_output_files(raw_output_files: Any) -> dict[str, dict[str, Any]]:
|
||||
def _enrich_metadata(path: str, metadata: dict[str, Any]) -> dict[str, Any]:
|
||||
normalized = dict(metadata)
|
||||
if "extension" not in normalized:
|
||||
normalized["extension"] = Path(path).suffix.lower().lstrip(".")
|
||||
if "mimetype" not in normalized:
|
||||
from abx_dl.output_files import guess_mimetype
|
||||
from abx_dl.output_files import OutputManifest
|
||||
|
||||
guessed = guess_mimetype(path)
|
||||
if guessed:
|
||||
normalized["mimetype"] = guessed
|
||||
return normalized
|
||||
|
||||
if raw_output_files is None:
|
||||
return {}
|
||||
if isinstance(raw_output_files, str):
|
||||
try:
|
||||
raw_output_files = json.loads(raw_output_files)
|
||||
except json.JSONDecodeError:
|
||||
return {}
|
||||
if isinstance(raw_output_files, dict):
|
||||
normalized: dict[str, dict[str, Any]] = {}
|
||||
for path, metadata in raw_output_files.items():
|
||||
if not path:
|
||||
continue
|
||||
metadata_dict = dict(metadata) if isinstance(metadata, dict) else {}
|
||||
metadata_dict.pop("path", None)
|
||||
normalized[str(path)] = _enrich_metadata(str(path), metadata_dict)
|
||||
return normalized
|
||||
if isinstance(raw_output_files, (list, tuple, set)):
|
||||
normalized: dict[str, dict[str, Any]] = {}
|
||||
for item in raw_output_files:
|
||||
if isinstance(item, str):
|
||||
normalized[item] = _enrich_metadata(item, {})
|
||||
continue
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
path = str(item.get("path") or "").strip()
|
||||
if not path:
|
||||
continue
|
||||
normalized[path] = _enrich_metadata(
|
||||
path,
|
||||
{key: value for key, value in item.items() if key != "path" and value not in (None, "")},
|
||||
)
|
||||
return normalized
|
||||
return {}
|
||||
return OutputManifest.from_value(raw_output_files).as_mapping()
|
||||
|
||||
@staticmethod
|
||||
def _coerce_output_file_size(value: Any) -> int:
|
||||
@ -4450,12 +4305,12 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes):
|
||||
return len(self.output_file_paths())
|
||||
|
||||
def output_size_from_files(self) -> int:
|
||||
return sum(self._coerce_output_file_size(metadata.get("size")) for metadata in self.output_file_map().values())
|
||||
from abx_dl.output_files import OutputManifest
|
||||
|
||||
return OutputManifest.from_value(self.output_files).total_size
|
||||
|
||||
def update_output_metadata_from_filesystem(self, snapshot_dir: Path | None = None, save: bool = True) -> bool:
|
||||
from collections import defaultdict
|
||||
|
||||
from abx_dl.output_files import guess_mimetype
|
||||
from abx_dl.output_files import OutputManifest, output_file_from_path
|
||||
|
||||
if self.plugin == "title":
|
||||
return False
|
||||
@ -4463,28 +4318,17 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes):
|
||||
snapshot_dir = Path(snapshot_dir or self.snapshot.output_dir)
|
||||
exclude_names = {"stdout.log", "stderr.log", "process.pid", "hook.pid", "listener.pid"}
|
||||
output_files: dict[str, dict[str, Any]] = {}
|
||||
mime_sizes: dict[str, int] = defaultdict(int)
|
||||
total_size = 0
|
||||
|
||||
def add_file(file_path: Path, rel_path: str, *, root_relative: bool = False) -> None:
|
||||
nonlocal total_size
|
||||
try:
|
||||
if not file_path.is_file() or file_path.name in exclude_names:
|
||||
return
|
||||
stat = file_path.stat()
|
||||
except OSError:
|
||||
return
|
||||
mime_type = guess_mimetype(file_path) or "application/octet-stream"
|
||||
metadata = {
|
||||
"extension": file_path.suffix.lower().lstrip("."),
|
||||
"mimetype": mime_type,
|
||||
"size": stat.st_size,
|
||||
}
|
||||
metadata = output_file_from_path(file_path, relative_to=file_path.parent).model_dump(exclude={"path"})
|
||||
if root_relative:
|
||||
metadata["root_relative"] = True
|
||||
output_files[rel_path] = metadata
|
||||
mime_sizes[mime_type] += stat.st_size
|
||||
total_size += stat.st_size
|
||||
|
||||
for raw_line in str(self.output_str or "").splitlines():
|
||||
raw_output = raw_line.strip().lstrip("/")
|
||||
@ -4506,16 +4350,14 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes):
|
||||
|
||||
plugin_dir = snapshot_dir / self.plugin
|
||||
if not output_files and plugin_dir.is_dir():
|
||||
for file_path in plugin_dir.rglob("*"):
|
||||
if not file_path.is_file() or ".hooks" in file_path.parts:
|
||||
continue
|
||||
add_file(file_path, str(file_path.relative_to(plugin_dir)))
|
||||
output_files = OutputManifest.scan(plugin_dir, containment_root=snapshot_dir).as_mapping()
|
||||
|
||||
if not output_files:
|
||||
return False
|
||||
|
||||
sorted_mimes = sorted(mime_sizes.items(), key=lambda item: item[1], reverse=True)
|
||||
output_mimetypes = ",".join(mime for mime, _ in sorted_mimes)
|
||||
manifest = OutputManifest.from_value(output_files)
|
||||
total_size = manifest.total_size
|
||||
output_mimetypes = ",".join(manifest.mimetypes)
|
||||
if self.output_files == output_files and self.output_size == total_size and self.output_mimetypes == output_mimetypes:
|
||||
return False
|
||||
|
||||
@ -4794,163 +4636,6 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes):
|
||||
def save_search_index(self):
|
||||
pass
|
||||
|
||||
def update_from_output(self):
|
||||
"""
|
||||
Update this ArchiveResult from filesystem logs and output files.
|
||||
|
||||
Used for Snapshot cleanup / orphan recovery when a hook's output exists
|
||||
on disk but the projector did not finalize the row in the database.
|
||||
|
||||
Updates:
|
||||
- status, output_str, output_json from ArchiveResult JSONL record
|
||||
- output_files, output_size, output_mimetypes by walking filesystem
|
||||
- end_ts, cmd, cmd_version, binary FK
|
||||
- Processes side-effect records (Snapshot, Tag, etc.) via process_hook_records()
|
||||
"""
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
from abx_dl.output_files import guess_mimetype
|
||||
from django.utils import timezone
|
||||
|
||||
from archivebox.machine.models import Process
|
||||
from archivebox.plugins.hooks import extract_records_from_process, process_hook_records
|
||||
|
||||
plugin_dir = Path(self.pwd) if self.pwd else None
|
||||
if not plugin_dir or not plugin_dir.exists():
|
||||
self.status = self.StatusChoices.FAILED
|
||||
self.output_str = "Output directory not found"
|
||||
self.end_ts = timezone.now()
|
||||
self.save()
|
||||
return
|
||||
|
||||
records = []
|
||||
process = self.process_record
|
||||
if process:
|
||||
records = extract_records_from_process(process)
|
||||
|
||||
if not records:
|
||||
stdout_file = plugin_dir / "stdout.log"
|
||||
stdout = stdout_file.read_text(errors="replace") if stdout_file.exists() else ""
|
||||
records = Process.parse_records_from_text(stdout)
|
||||
|
||||
# Find ArchiveResult record and update status/output from it
|
||||
ar_records = [r for r in records if r.get("type") == "ArchiveResult"]
|
||||
if ar_records:
|
||||
hook_data = ar_records[0]
|
||||
|
||||
# Update status
|
||||
status_map = {
|
||||
"succeeded": self.StatusChoices.SUCCEEDED,
|
||||
"failed": self.StatusChoices.FAILED,
|
||||
"skipped": self.StatusChoices.SKIPPED,
|
||||
"noresults": self.StatusChoices.NORESULTS,
|
||||
}
|
||||
self.status = status_map.get(hook_data.get("status", "failed"), self.StatusChoices.FAILED)
|
||||
|
||||
# Update output fields
|
||||
self.output_str = hook_data.get("output_str") or hook_data.get("output") or ""
|
||||
self.output_json = hook_data.get("output_json")
|
||||
|
||||
# Update cmd fields
|
||||
if hook_data.get("cmd"):
|
||||
if process:
|
||||
process.cmd = hook_data["cmd"]
|
||||
process.save()
|
||||
self._set_binary_from_cmd(hook_data["cmd"])
|
||||
# Note: cmd_version is derived from binary.version, not stored on Process
|
||||
else:
|
||||
# No ArchiveResult record: treat background hooks or clean exits as skipped
|
||||
is_background = False
|
||||
try:
|
||||
from archivebox.plugins.hooks import is_background_hook
|
||||
|
||||
is_background = bool(self.hook_name and is_background_hook(self.hook_name))
|
||||
except (ImportError, TypeError, ValueError):
|
||||
is_background = False
|
||||
|
||||
if is_background or (process and process.exit_code == 0):
|
||||
self.status = self.StatusChoices.SKIPPED
|
||||
self.output_str = "Hook did not output ArchiveResult record"
|
||||
else:
|
||||
self.status = self.StatusChoices.FAILED
|
||||
self.output_str = "Hook did not output ArchiveResult record"
|
||||
|
||||
# Walk filesystem and populate output_files, output_size, output_mimetypes
|
||||
exclude_names = {"stdout.log", "stderr.log", "process.pid", "hook.pid", "listener.pid"}
|
||||
mime_sizes = defaultdict(int)
|
||||
total_size = 0
|
||||
output_files = {}
|
||||
|
||||
for file_path in plugin_dir.rglob("*"):
|
||||
if not file_path.is_file():
|
||||
continue
|
||||
if ".hooks" in file_path.parts:
|
||||
continue
|
||||
if file_path.name in exclude_names:
|
||||
continue
|
||||
|
||||
try:
|
||||
stat = file_path.stat()
|
||||
mime_type = guess_mimetype(file_path) or "application/octet-stream"
|
||||
|
||||
relative_path = str(file_path.relative_to(plugin_dir))
|
||||
output_files[relative_path] = {
|
||||
"extension": file_path.suffix.lower().lstrip("."),
|
||||
"mimetype": mime_type,
|
||||
"size": stat.st_size,
|
||||
}
|
||||
mime_sizes[mime_type] += stat.st_size
|
||||
total_size += stat.st_size
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
self.output_files = output_files
|
||||
self.output_size = total_size
|
||||
sorted_mimes = sorted(mime_sizes.items(), key=lambda x: x[1], reverse=True)
|
||||
self.output_mimetypes = ",".join(mime for mime, _ in sorted_mimes)
|
||||
|
||||
# Update timestamps
|
||||
self.end_ts = timezone.now()
|
||||
|
||||
self.save()
|
||||
|
||||
# Process side-effect records (filter Snapshots for depth/URL)
|
||||
filtered_records = []
|
||||
for record in records:
|
||||
record_type = record.get("type")
|
||||
|
||||
# Skip ArchiveResult records (already processed above)
|
||||
if record_type == "ArchiveResult":
|
||||
continue
|
||||
|
||||
# Filter Snapshot records for depth/URL constraints
|
||||
if record_type == "Snapshot":
|
||||
url = record.get("url")
|
||||
if not url:
|
||||
continue
|
||||
|
||||
depth = record.get("depth", self.snapshot.depth + 1)
|
||||
if depth > self.snapshot.crawl.max_depth:
|
||||
continue
|
||||
|
||||
if not self._url_passes_filters(url):
|
||||
continue
|
||||
|
||||
filtered_records.append(record)
|
||||
|
||||
# Process filtered records with unified dispatcher
|
||||
overrides = {
|
||||
"snapshot": self.snapshot,
|
||||
"crawl": self.snapshot.crawl,
|
||||
"created_by_id": self.created_by.pk,
|
||||
}
|
||||
process_hook_records(filtered_records, overrides=overrides)
|
||||
|
||||
# Cleanup PID files (keep logs even if empty so they can be tailed)
|
||||
pid_file = plugin_dir / "hook.pid"
|
||||
pid_file.unlink(missing_ok=True)
|
||||
|
||||
def _set_binary_from_cmd(self, cmd: list) -> None:
|
||||
"""
|
||||
Find Binary for command and set binary FK.
|
||||
@ -5004,12 +4689,3 @@ class ArchiveResult(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithNotes):
|
||||
def output_dir(self) -> Path:
|
||||
"""Get the output directory for this plugin's results."""
|
||||
return Path(self.snapshot.output_dir) / self.plugin
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# State Machine Registration
|
||||
# =============================================================================
|
||||
|
||||
# Manually register state machines with python-statemachine registry
|
||||
# (normally auto-discovered from statemachines.py, but we define them here for clarity)
|
||||
registry.register(SnapshotMachine)
|
||||
|
||||
@ -13,7 +13,7 @@ def _is_signal_interrupted_exit(exit_code: int | None) -> bool:
|
||||
def recover_orchestrator_state(*, include_chrome: bool = False, crawl_id: str | None = None) -> dict[str, int]:
|
||||
from archivebox.crawls.models import Crawl
|
||||
from archivebox.core.models import ArchiveResult, Snapshot
|
||||
from archivebox.services.archive_result_service import _collect_output_metadata
|
||||
from abx_dl.output_files import OutputManifest
|
||||
from archivebox.machine.models import Process
|
||||
from django.core.exceptions import ValidationError
|
||||
from django.db.models import Exists, OuterRef, Q, Subquery, Value
|
||||
@ -200,7 +200,10 @@ def recover_orchestrator_state(*, include_chrome: bool = False, crawl_id: str |
|
||||
# A runner can die after the hook Process exits but before the
|
||||
# ProcessCompletedEvent projector links/finalizes ArchiveResult.
|
||||
# Reconstruct the plugin row from its newest durable Process row.
|
||||
output_files, output_size, output_mimetypes = _collect_output_metadata(plugin_dir)
|
||||
manifest = OutputManifest.scan(plugin_dir, containment_root=snapshot.output_dir)
|
||||
output_files = manifest.as_mapping()
|
||||
output_size = manifest.total_size
|
||||
output_mimetypes = ",".join(manifest.mimetypes)
|
||||
emitted_records = [
|
||||
record
|
||||
for record in Process.parse_records_from_text(process.stdout or "")
|
||||
@ -287,7 +290,7 @@ def recover_orchestrator_state(*, include_chrome: bool = False, crawl_id: str |
|
||||
|
||||
# Broken lock repair: STARTED + retry_at=NULL is an orphaned ownership
|
||||
# lease. Recovery only unlocks scheduling; the runner owns any subsequent
|
||||
# state-machine transition, including sealing rows whose children/results
|
||||
# lifecycle transition, including sealing rows whose children/results
|
||||
# are already final.
|
||||
recoverable_started_crawls = Crawl.objects.filter(status=Crawl.StatusChoices.STARTED).filter(
|
||||
Q(retry_at__isnull=True) | Q(retry_at__gt=now),
|
||||
|
||||
@ -68,14 +68,9 @@ _MEDIA_FILE_EXTS = {
|
||||
|
||||
|
||||
def _normalize_output_files(output_files: Any) -> dict[str, dict[str, Any]]:
|
||||
if isinstance(output_files, dict):
|
||||
normalized: dict[str, dict[str, Any]] = {}
|
||||
for path, metadata in output_files.items():
|
||||
if not path:
|
||||
continue
|
||||
normalized[str(path)] = dict(metadata) if isinstance(metadata, dict) else {}
|
||||
return normalized
|
||||
return {}
|
||||
from abx_dl.output_files import OutputManifest
|
||||
|
||||
return OutputManifest.from_value(output_files).as_mapping()
|
||||
|
||||
|
||||
def _snapshot_id(value: Any) -> Any:
|
||||
|
||||
@ -5,11 +5,3 @@ class CrawlsConfig(AppConfig):
|
||||
default_auto_field = "django.db.models.BigAutoField"
|
||||
name = "archivebox.crawls"
|
||||
label = "crawls"
|
||||
|
||||
def ready(self):
|
||||
"""Import models to register state machines with the registry"""
|
||||
import sys
|
||||
|
||||
# Skip during makemigrations to avoid premature state machine access
|
||||
if "makemigrations" not in sys.argv:
|
||||
from archivebox.crawls.models import CrawlMachine # noqa: F401
|
||||
|
||||
@ -19,7 +19,6 @@ from django.core.validators import MaxValueValidator, MinValueValidator
|
||||
from django.conf import settings
|
||||
from django.urls import reverse_lazy
|
||||
from django.utils import timezone
|
||||
from statemachine import State, registry
|
||||
from archivebox.config.common import rprint as print
|
||||
from archivebox.core.permissions import PERMISSIONS_VALUES, normalize_permissions
|
||||
|
||||
@ -32,7 +31,7 @@ from archivebox.base_models.models import (
|
||||
ModelWithHealthStats,
|
||||
get_or_create_system_user_pk,
|
||||
)
|
||||
from archivebox.workers.models import RETRY_AT_MAX, ModelWithStateMachine, BaseStateMachine
|
||||
from archivebox.workers.models import ModelWithQueue
|
||||
from archivebox.crawls.schedule_util import next_run_for_schedule, validate_schedule
|
||||
from archivebox.misc.util import parse_date, sanitize_html_text, validate_url, validate_url_length
|
||||
|
||||
@ -129,7 +128,7 @@ class CrawlSchedule(ModelWithUUID, ModelWithNotes):
|
||||
)
|
||||
|
||||
|
||||
class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWithHealthStats, ModelWithStateMachine):
|
||||
class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWithHealthStats, ModelWithQueue):
|
||||
id = CompactUUIDField(primary_key=True, default=uuid7, editable=False, unique=True)
|
||||
created_at = models.DateTimeField(default=timezone.now, db_index=True)
|
||||
created_by = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, default=get_or_create_system_user_pk, null=False)
|
||||
@ -158,16 +157,19 @@ class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWith
|
||||
notes = models.TextField(blank=True, null=False, default="")
|
||||
schedule = models.ForeignKey(CrawlSchedule, on_delete=models.SET_NULL, null=True, blank=True, editable=True)
|
||||
|
||||
status = ModelWithStateMachine.StatusField(
|
||||
choices=ModelWithStateMachine.StatusChoices,
|
||||
default=ModelWithStateMachine.StatusChoices.QUEUED,
|
||||
status = ModelWithQueue.StatusField(
|
||||
choices=ModelWithQueue.StatusChoices,
|
||||
default=ModelWithQueue.StatusChoices.QUEUED,
|
||||
)
|
||||
retry_at = ModelWithStateMachine.RetryAtField(default=timezone.now)
|
||||
retry_at = ModelWithQueue.RetryAtField(default=timezone.now)
|
||||
|
||||
state_machine_name = "archivebox.crawls.models.CrawlMachine"
|
||||
retry_at_field_name = "retry_at"
|
||||
state_field_name = "status"
|
||||
StatusChoices = ModelWithStateMachine.StatusChoices
|
||||
StatusChoices = ModelWithQueue.StatusChoices
|
||||
INITIAL_STATE = StatusChoices.QUEUED
|
||||
ACTIVE_STATE = StatusChoices.STARTED
|
||||
FINAL_STATES = (StatusChoices.SEALED,)
|
||||
FINAL_OR_ACTIVE_STATES = (*FINAL_STATES, ACTIVE_STATE)
|
||||
active_state = StatusChoices.STARTED
|
||||
delete_after_final_statuses = (StatusChoices.SEALED,)
|
||||
RUNNABLE_STATES = (StatusChoices.QUEUED, StatusChoices.STARTED)
|
||||
@ -177,17 +179,12 @@ class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWith
|
||||
|
||||
snapshot_set: models.Manager["Snapshot"]
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@property
|
||||
def sm(self) -> "CrawlMachine": ...
|
||||
|
||||
class Meta(
|
||||
ModelWithDeleteAfter.Meta,
|
||||
ModelWithOutputDir.Meta,
|
||||
ModelWithConfig.Meta,
|
||||
ModelWithHealthStats.Meta,
|
||||
ModelWithStateMachine.Meta,
|
||||
ModelWithQueue.Meta,
|
||||
):
|
||||
app_label = "crawls"
|
||||
verbose_name = "Crawl"
|
||||
@ -255,7 +252,7 @@ class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWith
|
||||
|
||||
now = timezone.now()
|
||||
# Cancellation seals the Crawl first, then lets the runner seal each
|
||||
# child Snapshot through its own state machine. Active children that
|
||||
# child Snapshot through its own lifecycle. Active children that
|
||||
# are already due need no write; the runner will claim them as-is.
|
||||
active_children = self.snapshot_set.filter(
|
||||
status__in=Snapshot.OPEN_STATES,
|
||||
@ -1297,7 +1294,7 @@ class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWith
|
||||
return created_snapshots
|
||||
|
||||
def install_declared_binaries(self, binary_names: set[str], machine=None) -> None:
|
||||
"""Install crawl-declared binaries through their unified state machine."""
|
||||
"""Install crawl-declared binaries through their unified lifecycle."""
|
||||
from archivebox.crawls.locks import binary_lifecycle_lock
|
||||
from archivebox.machine.models import Binary, Machine
|
||||
|
||||
@ -1313,7 +1310,7 @@ class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWith
|
||||
continue
|
||||
binary.update_and_requeue(retry_at=timezone.now())
|
||||
binary.refresh_from_db()
|
||||
binary.tick_claimed(lock_seconds=600)
|
||||
binary.install_claimed(lock_seconds=600)
|
||||
|
||||
unresolved_binaries = list(
|
||||
Binary.objects.filter(
|
||||
@ -1333,133 +1330,6 @@ class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWith
|
||||
f"Crawl dependencies failed to install before continuing: {binary_details}",
|
||||
)
|
||||
|
||||
def run(self) -> "Snapshot | None":
|
||||
"""
|
||||
Execute this Crawl: run hooks, process JSONL, create snapshots.
|
||||
|
||||
Called by the state machine when entering the 'started' state.
|
||||
|
||||
Returns:
|
||||
The root Snapshot for this crawl, or None for system crawls that don't create snapshots
|
||||
"""
|
||||
import time
|
||||
from archivebox.plugins.hooks import run_hook, discover_hooks, process_hook_records
|
||||
from archivebox.config.common import get_config
|
||||
from archivebox.machine.models import Machine
|
||||
|
||||
def get_runtime_config():
|
||||
return get_config(crawl=self).for_crawl_runtime(
|
||||
crawl=self,
|
||||
persona=persona,
|
||||
runtime_overrides=persona_runtime_overrides,
|
||||
)
|
||||
|
||||
system_task = self.get_system_task()
|
||||
if system_task == "archivebox://update":
|
||||
from archivebox.cli.archivebox_update import process_all_db_snapshots
|
||||
|
||||
process_all_db_snapshots()
|
||||
return None
|
||||
|
||||
machine = Machine.current()
|
||||
declared_binary_names: set[str] = set()
|
||||
persona_runtime_overrides: dict[str, str] = {}
|
||||
persona = self.resolve_persona()
|
||||
if persona:
|
||||
base_runtime_config = get_config(crawl=self, persona=persona)
|
||||
chrome_binary = str(base_runtime_config.get("CHROME_BINARY") or "")
|
||||
persona_runtime_overrides = persona.prepare_runtime_for_crawl(
|
||||
crawl=self,
|
||||
chrome_binary=chrome_binary,
|
||||
)
|
||||
|
||||
def run_crawl_hook(hook: Path) -> set[str]:
|
||||
primary_url = next(
|
||||
(line.strip() for line in self.urls.splitlines() if line.strip()),
|
||||
self.urls.strip(),
|
||||
)
|
||||
|
||||
hook_start = time.time()
|
||||
plugin_name = hook.parent.name
|
||||
output_dir = self.output_dir / plugin_name
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
process = run_hook(
|
||||
hook,
|
||||
output_dir=output_dir,
|
||||
config=get_runtime_config(),
|
||||
crawl_id=str(self.id),
|
||||
source_url=self.urls,
|
||||
url=primary_url,
|
||||
snapshot_id=str(self.id),
|
||||
)
|
||||
hook_elapsed = time.time() - hook_start
|
||||
if hook_elapsed > 0.5:
|
||||
print(f"[yellow]⏱️ Hook {hook.name} took {hook_elapsed:.2f}s[/yellow]")
|
||||
|
||||
if process.status == process.StatusChoices.RUNNING:
|
||||
if process.poll() is None:
|
||||
return set()
|
||||
|
||||
from archivebox.plugins.hooks import extract_records_from_process
|
||||
|
||||
records = []
|
||||
# A hook can exit before its completed Process metadata is visible.
|
||||
# Give successful hooks a brief chance to flush JSONL stdout into
|
||||
# the Process row before downstream hooks.
|
||||
for delay in (0.0, 0.05, 0.1, 0.25, 0.5):
|
||||
if delay:
|
||||
time.sleep(delay)
|
||||
records = extract_records_from_process(process)
|
||||
if records:
|
||||
break
|
||||
if records:
|
||||
print(f"[cyan]📝 Processing {len(records)} records from {hook.name}[/cyan]")
|
||||
for record in records[:3]:
|
||||
print(f" Record: type={record.get('type')}, keys={list(record.keys())[:5]}")
|
||||
if system_task:
|
||||
records = [record for record in records if record.get("type") in ("BinaryRequest", "Binary")]
|
||||
overrides = {"crawl": self}
|
||||
stats = process_hook_records(records, overrides=overrides)
|
||||
if stats:
|
||||
print(f"[green]✓ Created: {stats}[/green]")
|
||||
|
||||
hook_binary_names = {
|
||||
str(record.get("name")).strip()
|
||||
for record in records
|
||||
if record.get("type") in ("BinaryRequest", "Binary") and record.get("name")
|
||||
}
|
||||
hook_binary_names.discard("")
|
||||
if hook_binary_names:
|
||||
declared_binary_names.update(hook_binary_names)
|
||||
return hook_binary_names
|
||||
|
||||
hooks = discover_hooks("Crawl", config=get_runtime_config())
|
||||
|
||||
for hook in hooks:
|
||||
hook_binary_names = run_crawl_hook(hook)
|
||||
if hook_binary_names:
|
||||
self.install_declared_binaries(hook_binary_names, machine=machine)
|
||||
|
||||
# Safety check: don't create snapshots if any crawl-declared dependency
|
||||
# is still unresolved after all crawl hooks have run.
|
||||
self.install_declared_binaries(declared_binary_names, machine=machine)
|
||||
|
||||
# Create snapshots from all URLs in self.urls
|
||||
if system_task:
|
||||
leaked_snapshots = self.snapshot_set.all()
|
||||
if leaked_snapshots.exists():
|
||||
leaked_count = leaked_snapshots.count()
|
||||
leaked_snapshots.delete()
|
||||
print(f"[yellow]⚠️ Removed {leaked_count} leaked snapshot(s) created during system crawl {system_task}[/yellow]")
|
||||
return None
|
||||
|
||||
self.create_snapshots_from_urls()
|
||||
|
||||
# Return first snapshot for this crawl (newly created or existing)
|
||||
# This ensures the crawl doesn't seal if snapshots exist, even if they weren't just created
|
||||
return self.snapshot_set.first()
|
||||
|
||||
def is_finished(self) -> bool:
|
||||
"""Check if crawl is finished (all snapshots sealed or no snapshots exist)."""
|
||||
from archivebox.core.models import Snapshot
|
||||
@ -1483,11 +1353,64 @@ class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWith
|
||||
|
||||
return True
|
||||
|
||||
def cleanup(self):
|
||||
"""Clean up background hooks and run on_CrawlEnd hooks."""
|
||||
from archivebox.plugins.hooks import run_hook, discover_hooks
|
||||
def can_start(self) -> bool:
|
||||
return bool(self.urls and self.get_urls_list())
|
||||
|
||||
# Clean up .pid files from output directory
|
||||
def has_finished_snapshots(self) -> bool:
|
||||
from archivebox.core.models import Snapshot
|
||||
|
||||
snapshots = self.snapshot_set.all()
|
||||
return snapshots.exists() and not snapshots.exclude(status=Snapshot.StatusChoices.SEALED).exists()
|
||||
|
||||
def mark_started(self) -> bool:
|
||||
now = timezone.now()
|
||||
updated = self.safe_update(
|
||||
{
|
||||
"status": self.StatusChoices.STARTED,
|
||||
"retry_at": now + timedelta(seconds=2),
|
||||
},
|
||||
extra_filter={"status": self.StatusChoices.QUEUED},
|
||||
)
|
||||
return updated
|
||||
|
||||
def seal(self) -> bool:
|
||||
"""Finalize a runner-owned Crawl without dispatching hooks directly."""
|
||||
now = timezone.now()
|
||||
updated = self.safe_update(
|
||||
{
|
||||
"status": self.StatusChoices.SEALED,
|
||||
"retry_at": None,
|
||||
"modified_at": now,
|
||||
},
|
||||
refresh=False,
|
||||
extra_filter={"status__in": (*self.RUNNABLE_STATES, self.StatusChoices.SEALED)},
|
||||
)
|
||||
if not updated:
|
||||
self.refresh_from_db()
|
||||
return False
|
||||
self.status = self.StatusChoices.SEALED
|
||||
self.retry_at = None
|
||||
self.modified_at = now
|
||||
self.schedule_child_snapshots_for_sealing()
|
||||
self.cleanup_runtime()
|
||||
return True
|
||||
|
||||
def advance_lifecycle(self) -> bool:
|
||||
"""Advance one explicit lifecycle step after the runner claims this row."""
|
||||
if self.status == self.StatusChoices.PAUSED:
|
||||
return False
|
||||
if self.status == self.StatusChoices.QUEUED:
|
||||
if self.has_finished_snapshots():
|
||||
return self.seal()
|
||||
if not self.can_start():
|
||||
return False
|
||||
return self.mark_started()
|
||||
if self.status == self.StatusChoices.STARTED and self.is_finished():
|
||||
return self.seal()
|
||||
return False
|
||||
|
||||
def cleanup_runtime(self) -> None:
|
||||
"""Remove runner-owned runtime artifacts after abx-dl cleanup hooks finish."""
|
||||
if self.output_dir.exists():
|
||||
for pid_file in self.output_dir.glob("**/*.pid"):
|
||||
pid_file.unlink(missing_ok=True)
|
||||
@ -1495,207 +1418,3 @@ class Crawl(ModelWithDeleteAfter, ModelWithOutputDir, ModelWithConfig, ModelWith
|
||||
persona = self.resolve_persona()
|
||||
if persona:
|
||||
persona.cleanup_runtime_for_crawl(self)
|
||||
|
||||
# Run on_CrawlEnd hooks
|
||||
from archivebox.config.common import get_config
|
||||
|
||||
config = get_config(crawl=self)
|
||||
|
||||
hooks = discover_hooks("CrawlEnd", config=config)
|
||||
|
||||
for hook in hooks:
|
||||
plugin_name = hook.parent.name
|
||||
output_dir = self.output_dir / plugin_name
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
process = run_hook(
|
||||
hook,
|
||||
output_dir=output_dir,
|
||||
config=config,
|
||||
crawl_id=str(self.id),
|
||||
source_url=self.urls, # Pass full newline-separated URLs
|
||||
)
|
||||
|
||||
# Log failures but don't block
|
||||
if process.exit_code != 0:
|
||||
print(f"[yellow]⚠️ CrawlEnd hook failed: {hook.name}[/yellow]")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# State Machines
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class CrawlMachine(BaseStateMachine):
|
||||
crawl: Crawl
|
||||
|
||||
"""
|
||||
State machine for managing Crawl lifecycle.
|
||||
|
||||
Hook Lifecycle:
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ QUEUED State │
|
||||
│ • Waiting for crawl to be ready (has URLs) │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓ tick() when can_start()
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ STARTED State → enter_started() │
|
||||
│ 1. crawl.run() │
|
||||
│ • discover_hooks('Crawl') → finds all crawl hooks │
|
||||
│ • For each hook: │
|
||||
│ - run_hook(script, output_dir, ...) │
|
||||
│ - Parse JSONL from hook output │
|
||||
│ - process_hook_records() → creates Snapshots │
|
||||
│ • create_snapshots_from_urls() → from self.urls field │
|
||||
│ │
|
||||
│ 2. Snapshots process independently with their own │
|
||||
│ state machines (see SnapshotMachine) │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓ tick() when is_finished()
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ SEALED State → enter_sealed() │
|
||||
│ • cleanup() → runs on_CrawlEnd hooks, kills background │
|
||||
│ • Set retry_at=None (no more processing) │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
"""
|
||||
|
||||
model_attr_name = "crawl"
|
||||
|
||||
# States
|
||||
queued = State(value=Crawl.StatusChoices.QUEUED, initial=True)
|
||||
started = State(value=Crawl.StatusChoices.STARTED)
|
||||
paused = State(value=Crawl.StatusChoices.PAUSED)
|
||||
sealed = State(value=Crawl.StatusChoices.SEALED, final=True)
|
||||
|
||||
# Tick Event (polled by workers)
|
||||
tick = (
|
||||
queued.to(sealed, cond="has_finished_snapshots")
|
||||
| queued.to.itself(unless="can_start")
|
||||
| queued.to(started, cond="can_start")
|
||||
| started.to(sealed, cond="is_finished")
|
||||
| paused.to.itself()
|
||||
)
|
||||
|
||||
# Manual event (triggered by last Snapshot sealing, or by direct
|
||||
# index-only/bg creation when every requested URL is rejected before any
|
||||
# Snapshot rows exist).
|
||||
seal = queued.to(sealed) | started.to(sealed) | paused.to(sealed)
|
||||
pause_requested = queued.to(paused) | started.to(paused)
|
||||
resume_requested = paused.to(queued)
|
||||
|
||||
def can_start(self) -> bool:
|
||||
if not self.crawl.urls:
|
||||
print(f"[red]⚠️ Crawl {self.crawl.id} cannot start: no URLs[/red]")
|
||||
return False
|
||||
urls_list = self.crawl.get_urls_list()
|
||||
if not urls_list:
|
||||
print(f"[red]⚠️ Crawl {self.crawl.id} cannot start: no valid URLs in urls field[/red]")
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_finished(self) -> bool:
|
||||
"""Check if all Snapshots for this crawl are finished."""
|
||||
return self.crawl.is_finished()
|
||||
|
||||
def has_finished_snapshots(self) -> bool:
|
||||
"""A queued crawl with only final Snapshot rows was interrupted before sealing."""
|
||||
from archivebox.core.models import Snapshot
|
||||
|
||||
snapshots = self.crawl.snapshot_set.all()
|
||||
return snapshots.exists() and not snapshots.exclude(status=Snapshot.StatusChoices.SEALED).exists()
|
||||
|
||||
@queued.enter
|
||||
def enter_queued(self):
|
||||
self.crawl.update_and_requeue(
|
||||
retry_at=timezone.now(),
|
||||
status=Crawl.StatusChoices.QUEUED,
|
||||
)
|
||||
|
||||
@started.enter
|
||||
def enter_started(self):
|
||||
import sys
|
||||
|
||||
print(f"[cyan]🔄 CrawlMachine.enter_started() - creating snapshots for {self.crawl.id}[/cyan]", file=sys.stderr)
|
||||
|
||||
try:
|
||||
# Run the crawl - runs hooks, processes JSONL, creates snapshots
|
||||
first_snapshot = self.crawl.run()
|
||||
|
||||
if first_snapshot:
|
||||
print(
|
||||
f"[cyan]🔄 Created {self.crawl.snapshot_set.count()} snapshot(s), first: {first_snapshot.url}[/cyan]",
|
||||
file=sys.stderr,
|
||||
)
|
||||
# Update status to STARTED
|
||||
# Set retry_at to near future so tick() can poll and check is_finished()
|
||||
self.crawl.update_and_requeue(
|
||||
retry_at=timezone.now() + timedelta(seconds=2),
|
||||
status=Crawl.StatusChoices.STARTED,
|
||||
)
|
||||
else:
|
||||
# No snapshots (system crawl that only runs setup hooks)
|
||||
print("[cyan]🔄 No snapshots created, sealing crawl immediately[/cyan]", file=sys.stderr)
|
||||
# Seal immediately since there's no work to do
|
||||
self.seal()
|
||||
|
||||
except Exception as e:
|
||||
print(f"[red]⚠️ Crawl {self.crawl.id} failed to start: {e}[/red]")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
raise
|
||||
|
||||
@paused.enter
|
||||
def enter_paused(self):
|
||||
paused = self.crawl.safe_update(
|
||||
{
|
||||
"retry_at": RETRY_AT_MAX,
|
||||
"status": Crawl.StatusChoices.PAUSED,
|
||||
},
|
||||
extra_filter={"status__in": Crawl.RUNNABLE_STATES},
|
||||
)
|
||||
if paused:
|
||||
self.crawl.schedule_child_snapshots_for_pause()
|
||||
|
||||
@sealed.enter
|
||||
def enter_sealed(self):
|
||||
now = timezone.now()
|
||||
self.crawl.status = Crawl.StatusChoices.SEALED
|
||||
self.crawl.retry_at = None
|
||||
# Guard: never seal a row that a concurrent writer flipped to PAUSED.
|
||||
# Sealing is idempotent (SEALED→SEALED is a no-op rewrite), so
|
||||
# status__in covers both the QUEUED/STARTED→SEALED transition and the
|
||||
# rare re-entry case.
|
||||
updated = self.crawl.safe_update(
|
||||
{
|
||||
"status": Crawl.StatusChoices.SEALED,
|
||||
"retry_at": None,
|
||||
"modified_at": now,
|
||||
},
|
||||
refresh=False,
|
||||
extra_filter={
|
||||
"status__in": [
|
||||
Crawl.StatusChoices.QUEUED,
|
||||
Crawl.StatusChoices.STARTED,
|
||||
Crawl.StatusChoices.SEALED,
|
||||
],
|
||||
},
|
||||
)
|
||||
if not updated:
|
||||
self.crawl.refresh_from_db()
|
||||
return
|
||||
self.crawl.modified_at = now
|
||||
|
||||
self.crawl.schedule_child_snapshots_for_sealing()
|
||||
# Clean up background hooks and run on_CrawlEnd hooks after the final
|
||||
# state is visible so cleanup projectors cannot resurrect the crawl.
|
||||
self.crawl.cleanup()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Register State Machines
|
||||
# =============================================================================
|
||||
|
||||
# Manually register state machines with python-statemachine registry
|
||||
# (normally auto-discovered from statemachines.py, but we define them here for clarity)
|
||||
registry.register(CrawlMachine)
|
||||
|
||||
@ -10,14 +10,6 @@ class MachineConfig(AppConfig):
|
||||
label = "machine" # Explicit label for migrations
|
||||
verbose_name = "Machine Info"
|
||||
|
||||
def ready(self):
|
||||
"""Import models to register state machines with the registry"""
|
||||
import sys
|
||||
|
||||
# Skip during makemigrations to avoid premature state machine access
|
||||
if "makemigrations" not in sys.argv:
|
||||
from archivebox.machine import models # noqa: F401
|
||||
|
||||
|
||||
def register_admin(admin_site):
|
||||
from archivebox.machine.admin import register_admin
|
||||
|
||||
@ -12,8 +12,6 @@ from archivebox.uuid_compat import CompactUUIDField, uuid7
|
||||
from datetime import timedelta, datetime
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
from statemachine import State, registry
|
||||
|
||||
from django.db import IntegrityError, transaction
|
||||
from django.db import models
|
||||
from django.db.models import Q, QuerySet
|
||||
@ -23,7 +21,7 @@ from django.utils.functional import cached_property
|
||||
from archivebox.config import CONSTANTS
|
||||
from archivebox.config.common import rprint
|
||||
from archivebox.base_models.models import ModelWithDeleteAfter, ModelWithHealthStats, normalize_config_json_values
|
||||
from archivebox.workers.models import BaseStateMachine, ModelWithStateMachine
|
||||
from archivebox.workers.models import ModelWithQueue
|
||||
from .detect import get_host_guid, get_os_info, get_vm_info, get_host_network, get_host_stats
|
||||
|
||||
_psutil: Any | None = None
|
||||
@ -515,11 +513,11 @@ class BinaryManager(models.Manager):
|
||||
)
|
||||
|
||||
|
||||
class Binary(ModelWithHealthStats, ModelWithStateMachine):
|
||||
class Binary(ModelWithHealthStats, ModelWithQueue):
|
||||
"""
|
||||
Tracks a binary on a specific machine.
|
||||
|
||||
Simple state machine with 2 states:
|
||||
Simple queue lifecycle with 2 states:
|
||||
- queued: Binary needs to be installed
|
||||
- installed: Binary installed successfully (abspath, version, sha256 populated)
|
||||
|
||||
@ -567,8 +565,8 @@ class Binary(ModelWithHealthStats, ModelWithStateMachine):
|
||||
sha256 = models.CharField(max_length=64, default="", null=False, blank=True)
|
||||
|
||||
# State machine fields
|
||||
status = ModelWithStateMachine.StatusField(choices=StatusChoices.choices, default=StatusChoices.QUEUED, max_length=16)
|
||||
retry_at = ModelWithStateMachine.RetryAtField(
|
||||
status = ModelWithQueue.StatusField(choices=StatusChoices.choices, default=StatusChoices.QUEUED, max_length=16)
|
||||
retry_at = ModelWithQueue.RetryAtField(
|
||||
default=timezone.now,
|
||||
help_text="When to retry this binary installation",
|
||||
)
|
||||
@ -579,18 +577,16 @@ class Binary(ModelWithHealthStats, ModelWithStateMachine):
|
||||
|
||||
machine_id: uuid.UUID
|
||||
|
||||
state_machine_name: str | None = "archivebox.machine.models.BinaryMachine"
|
||||
INITIAL_STATE = StatusChoices.QUEUED
|
||||
ACTIVE_STATE = StatusChoices.QUEUED
|
||||
FINAL_STATES = (StatusChoices.INSTALLED,)
|
||||
FINAL_OR_ACTIVE_STATES = (*FINAL_STATES, ACTIVE_STATE)
|
||||
active_state: str = StatusChoices.QUEUED
|
||||
warn_on_save_outside_runner = False
|
||||
|
||||
objects = BinaryManager() # pyright: ignore[reportIncompatibleVariableOverride]
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@property
|
||||
def sm(self) -> BinaryMachine: ...
|
||||
|
||||
class Meta(ModelWithHealthStats.Meta, ModelWithStateMachine.Meta):
|
||||
class Meta(ModelWithHealthStats.Meta, ModelWithQueue.Meta):
|
||||
app_label = "machine"
|
||||
verbose_name = "Binary"
|
||||
verbose_name_plural = "Binaries"
|
||||
@ -604,6 +600,10 @@ class Binary(ModelWithHealthStats, ModelWithStateMachine):
|
||||
"""A binary is valid if it has a resolved path and is marked installed."""
|
||||
return bool(self.abspath) and self.status == self.StatusChoices.INSTALLED
|
||||
|
||||
@property
|
||||
def can_install(self) -> bool:
|
||||
return bool(self.name and self.binproviders)
|
||||
|
||||
@cached_property
|
||||
def binary_info(self) -> dict:
|
||||
"""Return info about the binary."""
|
||||
@ -748,11 +748,43 @@ class Binary(ModelWithHealthStats, ModelWithStateMachine):
|
||||
|
||||
run_binary(str(self.id))
|
||||
|
||||
def install(self) -> bool:
|
||||
"""Run one synchronous installation attempt for a claimed Binary."""
|
||||
if self.status == self.StatusChoices.INSTALLED:
|
||||
return True
|
||||
if not self.can_install:
|
||||
return False
|
||||
|
||||
rprint(f"[cyan] 🔄 installing {self.name}[/cyan]", file=sys.stderr)
|
||||
self.run()
|
||||
self.refresh_from_db()
|
||||
if self.status != self.StatusChoices.INSTALLED:
|
||||
self.update_and_requeue(
|
||||
retry_at=timezone.now() + timedelta(seconds=300),
|
||||
status=self.StatusChoices.QUEUED,
|
||||
)
|
||||
self.increment_health_stats(success=False)
|
||||
raise RuntimeError(f"Binary {self.name} installation failed")
|
||||
|
||||
self.update_and_requeue(retry_at=None, status=self.StatusChoices.INSTALLED)
|
||||
self.increment_health_stats(success=True)
|
||||
return True
|
||||
|
||||
def advance_lifecycle(self) -> bool:
|
||||
"""Advance the explicit binary lifecycle after its queue row is claimed."""
|
||||
return self.install()
|
||||
|
||||
def install_claimed(self, *, lock_seconds: int = 600) -> bool:
|
||||
if not self.claim_processing_lock(lock_seconds=lock_seconds):
|
||||
return False
|
||||
self.refresh_from_db()
|
||||
return self.advance_lifecycle()
|
||||
|
||||
def cleanup(self):
|
||||
"""
|
||||
Clean up background binary installation hooks.
|
||||
|
||||
Called by state machine if needed (not typically used for binaries
|
||||
Called after an installation attempt if needed (not typically used for binaries
|
||||
since installations are foreground, but included for consistency).
|
||||
"""
|
||||
|
||||
@ -953,7 +985,7 @@ class Process(ModelWithDeleteAfter, models.Model):
|
||||
One Process can optionally be associated with an ArchiveResult (via OneToOne),
|
||||
but Process can also exist standalone for internal operations.
|
||||
|
||||
Follows the unified state machine pattern:
|
||||
Follows the unified process lifecycle:
|
||||
- queued: Process ready to launch
|
||||
- running: Process actively executing
|
||||
- exited: Process completed (check exit_code for success/failure)
|
||||
@ -1138,7 +1170,6 @@ class Process(ModelWithDeleteAfter, models.Model):
|
||||
children: models.Manager[Process]
|
||||
archiveresult: ArchiveResult
|
||||
|
||||
state_machine_name: str = "archivebox.machine.models.ProcessMachine"
|
||||
delete_after_final_statuses = (StatusChoices.EXITED,)
|
||||
|
||||
objects = ProcessManager() # pyright: ignore[reportIncompatibleVariableOverride]
|
||||
@ -1303,8 +1334,8 @@ class Process(ModelWithDeleteAfter, models.Model):
|
||||
"""
|
||||
Compare-and-swap update for short Process scheduler writes.
|
||||
|
||||
Process is not a ModelWithStateMachine subclass yet, but its
|
||||
state-machine methods still need the same modified_at CAS behavior as
|
||||
Process is not a ModelWithQueue subclass, but its scheduler methods
|
||||
still need the same modified_at CAS behavior as
|
||||
Crawl/Snapshot/Binary without falling back to save().
|
||||
"""
|
||||
values = dict(update_fields)
|
||||
@ -2593,206 +2624,3 @@ class Process(ModelWithDeleteAfter, models.Model):
|
||||
if cleaned:
|
||||
rprint(f"[yellow]🧹 Cleaned up {cleaned} orphaned worker/hook process record(s)[/yellow]")
|
||||
return cleaned
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Binary State Machine
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class BinaryMachine(BaseStateMachine):
|
||||
"""
|
||||
State machine for managing Binary installation lifecycle.
|
||||
|
||||
Simple 2-state machine:
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ QUEUED State │
|
||||
│ • Binary needs to be installed │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓ tick() when can_install()
|
||||
↓ Synchronous installation during transition
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ INSTALLED State │
|
||||
│ • Binary installed (abspath, version, sha256 set) │
|
||||
│ • Health stats incremented │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
|
||||
If installation fails, Binary stays in QUEUED with retry_at bumped.
|
||||
"""
|
||||
|
||||
model_attr_name = "binary"
|
||||
binary: Binary
|
||||
|
||||
# States
|
||||
queued = State(value=Binary.StatusChoices.QUEUED, initial=True)
|
||||
installed = State(value=Binary.StatusChoices.INSTALLED, final=True)
|
||||
|
||||
# Tick Event - install happens during transition
|
||||
tick = queued.to.itself(unless="can_install") | queued.to(installed, cond="can_install", on="on_install")
|
||||
|
||||
def can_install(self) -> bool:
|
||||
"""Check if binary installation can start."""
|
||||
return bool(self.binary.name and self.binary.binproviders)
|
||||
|
||||
@queued.enter
|
||||
def enter_queued(self):
|
||||
"""Binary is queued for installation."""
|
||||
self.binary.update_and_requeue(
|
||||
retry_at=timezone.now(),
|
||||
status=Binary.StatusChoices.QUEUED,
|
||||
)
|
||||
|
||||
def on_install(self):
|
||||
"""Called during queued→installed transition. Runs installation synchronously."""
|
||||
import sys
|
||||
|
||||
rprint(f"[cyan] 🔄 BinaryMachine.on_install() - installing {self.binary.name}[/cyan]", file=sys.stderr)
|
||||
|
||||
# Run installation hooks (synchronous, updates abspath/version/sha256 and sets status)
|
||||
self.binary.run()
|
||||
|
||||
# Check if installation succeeded by looking at updated status
|
||||
# Note: Binary.run() updates self.binary.status internally but doesn't refresh our reference
|
||||
self.binary.refresh_from_db()
|
||||
|
||||
if self.binary.status != Binary.StatusChoices.INSTALLED:
|
||||
# Installation failed - abort transition, stay in queued
|
||||
rprint(f"[red] ❌ BinaryMachine - {self.binary.name} installation failed, retrying later[/red]", file=sys.stderr)
|
||||
|
||||
# Bump retry_at to try again later
|
||||
self.binary.update_and_requeue(
|
||||
retry_at=timezone.now() + timedelta(seconds=300), # Retry in 5 minutes
|
||||
status=Binary.StatusChoices.QUEUED, # Ensure we stay queued
|
||||
)
|
||||
|
||||
# Increment health stats for failure
|
||||
self.binary.increment_health_stats(success=False)
|
||||
|
||||
# Abort the transition - this will raise an exception and keep us in queued
|
||||
raise Exception(f"Binary {self.binary.name} installation failed")
|
||||
|
||||
rprint(f"[cyan] ✅ BinaryMachine - {self.binary.name} installed successfully[/cyan]", file=sys.stderr)
|
||||
|
||||
@installed.enter
|
||||
def enter_installed(self):
|
||||
"""Binary installed successfully."""
|
||||
self.binary.update_and_requeue(
|
||||
retry_at=None,
|
||||
status=Binary.StatusChoices.INSTALLED,
|
||||
)
|
||||
|
||||
# Increment health stats
|
||||
self.binary.increment_health_stats(success=True)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Process State Machine
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class ProcessMachine(BaseStateMachine):
|
||||
"""
|
||||
State machine for managing Process (OS subprocess) lifecycle.
|
||||
|
||||
Process Lifecycle:
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ QUEUED State │
|
||||
│ • Process ready to launch, waiting for resources │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓ tick() when can_start()
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ RUNNING State → enter_running() │
|
||||
│ 1. process.launch() │
|
||||
│ • Spawn subprocess with cmd, pwd, env, timeout │
|
||||
│ • Set pid, started_at │
|
||||
│ • Process runs in background or foreground │
|
||||
│ 2. Monitor process completion │
|
||||
│ • Check exit code when process completes │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓ tick() checks is_exited()
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ EXITED State │
|
||||
│ • Process completed (exit_code set) │
|
||||
│ • Health stats incremented │
|
||||
│ • stdout/stderr captured │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
|
||||
Note: This is a simpler state machine than ArchiveResult.
|
||||
Process is just about execution lifecycle. ArchiveResult handles
|
||||
the archival-specific logic (status, output parsing, etc.).
|
||||
"""
|
||||
|
||||
model_attr_name = "process"
|
||||
process: Process
|
||||
|
||||
# States
|
||||
queued = State(value=Process.StatusChoices.QUEUED, initial=True)
|
||||
running = State(value=Process.StatusChoices.RUNNING)
|
||||
exited = State(value=Process.StatusChoices.EXITED, final=True)
|
||||
|
||||
# Tick Event - transitions based on conditions
|
||||
tick = (
|
||||
queued.to.itself(unless="can_start")
|
||||
| queued.to(running, cond="can_start")
|
||||
| running.to.itself(unless="is_exited")
|
||||
| running.to(exited, cond="is_exited")
|
||||
)
|
||||
|
||||
# Additional events (for explicit control)
|
||||
launch = queued.to(running)
|
||||
kill = running.to(exited)
|
||||
|
||||
def can_start(self) -> bool:
|
||||
"""Check if process can start (has cmd and machine)."""
|
||||
return bool(self.process.cmd and self.process.machine)
|
||||
|
||||
def is_exited(self) -> bool:
|
||||
"""Check if process has exited (exit_code is set)."""
|
||||
return self.process.exit_code is not None
|
||||
|
||||
@queued.enter
|
||||
def enter_queued(self):
|
||||
"""Process is queued for execution."""
|
||||
self.process.update_and_requeue(
|
||||
retry_at=timezone.now(),
|
||||
status=Process.StatusChoices.QUEUED,
|
||||
)
|
||||
|
||||
@running.enter
|
||||
def enter_running(self):
|
||||
"""Start process execution."""
|
||||
# Lock the process while it runs
|
||||
self.process.update_and_requeue(
|
||||
retry_at=timezone.now() + timedelta(seconds=self.process.timeout),
|
||||
status=Process.StatusChoices.RUNNING,
|
||||
started_at=timezone.now(),
|
||||
)
|
||||
|
||||
# Launch the subprocess
|
||||
# NOTE: This is a placeholder - actual launch logic would
|
||||
# be implemented based on how hooks currently spawn processes
|
||||
# For now, Process is a data model that tracks execution metadata
|
||||
# The actual subprocess spawning is still handled by run_hook()
|
||||
|
||||
# Mark as immediately exited for now (until we refactor run_hook)
|
||||
# In the future, this would actually spawn the subprocess
|
||||
self.process.exit_code = 0 # Placeholder
|
||||
self.process.save()
|
||||
|
||||
@exited.enter
|
||||
def enter_exited(self):
|
||||
"""Process has exited."""
|
||||
self.process.update_and_requeue(
|
||||
retry_at=None,
|
||||
status=Process.StatusChoices.EXITED,
|
||||
ended_at=timezone.now(),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# State Machine Registration
|
||||
# =============================================================================
|
||||
|
||||
# Manually register state machines with python-statemachine registry
|
||||
registry.register(BinaryMachine)
|
||||
registry.register(ProcessMachine)
|
||||
|
||||
@ -1,12 +1,12 @@
|
||||
__package__ = "archivebox.plugins"
|
||||
|
||||
import json
|
||||
from collections.abc import Iterable
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import Any, Protocol, TypedDict
|
||||
|
||||
from abx_plugins import get_plugins_dir
|
||||
from abx_dl.catalog import PluginCatalog, PluginConfigResolver
|
||||
from django.utils.safestring import mark_safe
|
||||
|
||||
from archivebox.config.constants import CONSTANTS
|
||||
@ -29,18 +29,18 @@ USER_PLUGINS_DIR = CONSTANTS.USER_PLUGINS_DIR
|
||||
|
||||
|
||||
def iter_plugin_dirs() -> list[Path]:
|
||||
"""Iterate over all built-in and user plugin directories."""
|
||||
plugin_dirs: list[Path] = []
|
||||
"""Return the exact plugin directories exposed by the shared catalog."""
|
||||
return [plugin.path for plugin in get_plugin_catalog().values()]
|
||||
|
||||
for base_dir in (BUILTIN_PLUGINS_DIR, USER_PLUGINS_DIR):
|
||||
if not base_dir.exists():
|
||||
continue
|
||||
|
||||
for plugin_dir in base_dir.iterdir():
|
||||
if plugin_dir.is_dir() and not plugin_dir.name.startswith("_"):
|
||||
plugin_dirs.append(plugin_dir)
|
||||
@lru_cache(maxsize=1)
|
||||
def get_plugin_catalog() -> PluginCatalog:
|
||||
return PluginCatalog.discover(extra_plugin_dirs=[USER_PLUGINS_DIR], runtime="archivebox")
|
||||
|
||||
return plugin_dirs
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_plugin_config_resolver() -> PluginConfigResolver:
|
||||
return PluginConfigResolver(get_plugin_catalog())
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
@ -52,25 +52,11 @@ def get_plugins() -> list[str]:
|
||||
or a standardized templates/icon.html asset. This includes non-extractor
|
||||
plugins such as binary providers and shared base plugins.
|
||||
"""
|
||||
plugins = []
|
||||
|
||||
for plugin_dir in iter_plugin_dirs():
|
||||
has_hooks = any(plugin_dir.glob("on_*__*.*"))
|
||||
has_config = (plugin_dir / "config.json").exists()
|
||||
has_icon = (plugin_dir / "templates" / "icon.html").exists()
|
||||
if has_hooks or has_config or has_icon:
|
||||
plugins.append(plugin_dir.name)
|
||||
|
||||
return sorted(set(plugins))
|
||||
return sorted(get_plugin_catalog())
|
||||
|
||||
|
||||
def get_plugin_models():
|
||||
from abx_dl.models import discover_plugins
|
||||
|
||||
plugins = {}
|
||||
for base_dir in (BUILTIN_PLUGINS_DIR, USER_PLUGINS_DIR):
|
||||
plugins.update(discover_plugins(plugins_dir=base_dir, runtime="archivebox"))
|
||||
return plugins
|
||||
return get_plugin_catalog().plugins
|
||||
|
||||
|
||||
def get_plugin_name(plugin: str) -> str:
|
||||
@ -99,18 +85,7 @@ def get_enabled_plugins(config: ConfigLookup | None = None, **config_kwargs: Any
|
||||
|
||||
config = get_config(**config_kwargs)
|
||||
|
||||
enabled = []
|
||||
disabled = []
|
||||
for plugin in get_plugins():
|
||||
plugin_config = get_plugin_special_config(plugin, config)
|
||||
if plugin_config["enabled"]:
|
||||
enabled.append(plugin)
|
||||
else:
|
||||
disabled.append(plugin)
|
||||
|
||||
from abx_dl.models import filter_plugins
|
||||
|
||||
return list(filter_plugins(get_plugin_models(), enabled, include_providers=True, disabled_names=disabled))
|
||||
return get_plugin_config_resolver().enabled_plugin_names_from_flat(dict(config.items()))
|
||||
|
||||
|
||||
def discover_plugins_that_provide_interface(
|
||||
@ -128,14 +103,7 @@ def discover_plugins_that_provide_interface(
|
||||
|
||||
backends = {}
|
||||
|
||||
for base_dir in (BUILTIN_PLUGINS_DIR, USER_PLUGINS_DIR):
|
||||
if not base_dir.exists():
|
||||
continue
|
||||
|
||||
for plugin_dir in base_dir.iterdir():
|
||||
if not plugin_dir.is_dir():
|
||||
continue
|
||||
|
||||
for plugin_dir in iter_plugin_dirs():
|
||||
plugin_name = plugin_dir.name
|
||||
if plugin_prefix and not plugin_name.startswith(plugin_prefix):
|
||||
continue
|
||||
@ -196,33 +164,7 @@ def discover_plugin_configs() -> dict[str, dict[str, Any]]:
|
||||
schemas are plugin package metadata, not live user config; runtime values
|
||||
still come from env/db config at each callsite.
|
||||
"""
|
||||
configs = {}
|
||||
|
||||
for plugin_dir in iter_plugin_dirs():
|
||||
config_path = plugin_dir / "config.json"
|
||||
if not config_path.exists():
|
||||
continue
|
||||
|
||||
try:
|
||||
with open(config_path) as f:
|
||||
schema = json.load(f)
|
||||
|
||||
if not isinstance(schema, dict):
|
||||
continue
|
||||
if schema.get("type") != "object":
|
||||
continue
|
||||
if "properties" not in schema:
|
||||
continue
|
||||
|
||||
configs[plugin_dir.name] = schema
|
||||
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
import sys
|
||||
|
||||
print(f"Warning: Failed to load config.json from {plugin_dir.name}: {e}", file=sys.stderr)
|
||||
continue
|
||||
|
||||
return configs
|
||||
return get_plugin_config_resolver().schemas
|
||||
|
||||
|
||||
def get_plugin_special_config(plugin_name: str, config: ConfigLookup, _visited: set[str] | None = None) -> PluginSpecialConfig:
|
||||
@ -234,26 +176,7 @@ def get_plugin_special_config(plugin_name: str, config: ConfigLookup, _visited:
|
||||
- {PLUGIN}_TIMEOUT: Plugin-specific timeout (fallback to TIMEOUT, default 300)
|
||||
- {PLUGIN}_BINARY: Primary binary path (default to plugin_name)
|
||||
"""
|
||||
plugin_upper = plugin_name.upper()
|
||||
|
||||
enabled_key = f"{plugin_upper}_ENABLED"
|
||||
enabled = config.get(enabled_key)
|
||||
if enabled is None:
|
||||
enabled = True
|
||||
elif isinstance(enabled, str):
|
||||
enabled = enabled.lower() not in ("false", "0", "no", "")
|
||||
|
||||
timeout_key = f"{plugin_upper}_TIMEOUT"
|
||||
timeout = config.get(timeout_key) or config.get("TIMEOUT", 300)
|
||||
|
||||
binary_key = f"{plugin_upper}_BINARY"
|
||||
binary = config.get(binary_key, plugin_name)
|
||||
|
||||
return {
|
||||
"enabled": bool(enabled),
|
||||
"timeout": int(timeout),
|
||||
"binary": str(binary),
|
||||
}
|
||||
return get_plugin_config_resolver().runtime_settings(plugin_name, dict(config.items()))
|
||||
|
||||
|
||||
DEFAULT_TEMPLATES = {
|
||||
@ -295,10 +218,10 @@ def get_plugin_template(plugin: str, template_name: str, fallback: bool = True)
|
||||
if base_name in ("yt-dlp", "youtube-dl"):
|
||||
base_name = "ytdlp"
|
||||
|
||||
for plugin_dir in iter_plugin_dirs():
|
||||
if plugin_dir.name == base_name or plugin_dir.name.endswith(f"_{base_name}"):
|
||||
template_path = plugin_dir / "templates" / f"{template_name}.html"
|
||||
if template_path.exists():
|
||||
catalog = get_plugin_catalog()
|
||||
if base_name in catalog:
|
||||
template_path = catalog.template_path(base_name, template_name)
|
||||
if template_path is not None:
|
||||
return template_path.read_text()
|
||||
|
||||
if fallback:
|
||||
|
||||
@ -11,133 +11,19 @@ from django.utils.html import format_html
|
||||
|
||||
from archivebox.config import CONSTANTS_CONFIG
|
||||
from archivebox.config.common import ArchiveBoxConfig, get_config
|
||||
from archivebox.plugins.discovery import discover_plugin_configs, get_plugin_icon, get_plugins
|
||||
from archivebox.plugins.discovery import discover_plugin_configs, get_plugin_catalog, get_plugin_icon, get_plugins
|
||||
|
||||
|
||||
PLUGIN_CONFIG_FIELD_PREFIX = "plugin_config__"
|
||||
PLUGIN_GROUP_DEFINITIONS = (
|
||||
(
|
||||
"main_plugins",
|
||||
"Main",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
(
|
||||
"dom",
|
||||
"screenshot",
|
||||
"pdf",
|
||||
"singlefile",
|
||||
"wget",
|
||||
"archivedotorg",
|
||||
"chrome_mhtml",
|
||||
"archivewebpage",
|
||||
),
|
||||
),
|
||||
(
|
||||
"page_setup_plugins",
|
||||
"Page Setup",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
(
|
||||
"chrome",
|
||||
"infiniscroll",
|
||||
"modalcloser",
|
||||
"ublock",
|
||||
"istilldontcareaboutcookies",
|
||||
"twocaptcha",
|
||||
"claudechrome",
|
||||
),
|
||||
),
|
||||
(
|
||||
"media_plugins",
|
||||
"Media",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
(
|
||||
"staticfile",
|
||||
"responses",
|
||||
"chrome_screencast",
|
||||
"ytdlp",
|
||||
"gallerydl",
|
||||
"git",
|
||||
),
|
||||
),
|
||||
(
|
||||
"text_plugins",
|
||||
"Text",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
(
|
||||
"readability",
|
||||
"htmltotext",
|
||||
"defuddle",
|
||||
"forumdl",
|
||||
"mercury",
|
||||
"trafilatura",
|
||||
"liteparse",
|
||||
"opendataloader",
|
||||
"papersdl",
|
||||
),
|
||||
),
|
||||
(
|
||||
"metadata_plugins",
|
||||
"Metadata",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
(
|
||||
"title",
|
||||
"favicon",
|
||||
"headers",
|
||||
"redirects",
|
||||
"accessibility",
|
||||
"consolelog",
|
||||
"sslcerts",
|
||||
"dns",
|
||||
"seo",
|
||||
"hashes",
|
||||
),
|
||||
),
|
||||
(
|
||||
"postprocessing_plugins",
|
||||
"Postprocessing",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
(
|
||||
"parse_dom_outlinks",
|
||||
"parse_html_urls",
|
||||
"parse_jsonl_urls",
|
||||
"parse_netscape_urls",
|
||||
"parse_rss_urls",
|
||||
"parse_txt_urls",
|
||||
"claudecode",
|
||||
"claudecodecleanup",
|
||||
"claudecodeextract",
|
||||
),
|
||||
),
|
||||
PLUGIN_GROUPS = (
|
||||
("main", "main_plugins", "Main"),
|
||||
("page_setup", "page_setup_plugins", "Page Setup"),
|
||||
("media", "media_plugins", "Media"),
|
||||
("text", "text_plugins", "Text"),
|
||||
("metadata", "metadata_plugins", "Metadata"),
|
||||
("postprocessing", "postprocessing_plugins", "Postprocessing"),
|
||||
("other", "other_plugins", "Other"),
|
||||
)
|
||||
HIDDEN_PLUGIN_CONFIG_UI_PLUGINS = {
|
||||
"apt",
|
||||
"base",
|
||||
"bash",
|
||||
"brew",
|
||||
"cargo",
|
||||
"chromewebstore",
|
||||
"env",
|
||||
"media",
|
||||
"npm",
|
||||
"opencode",
|
||||
"pip",
|
||||
"puppeteer",
|
||||
"search_backend_ripgrep",
|
||||
"search_backend_sonic",
|
||||
"search_backend_sqlite",
|
||||
"ssl",
|
||||
}
|
||||
TIMEOUT_INPUT_PATTERN = r"(0|[1-9][0-9]*|[0-9]+(?:\.[0-9]+)?\s*(?:s|sec|secs|second|seconds|m|min|mins|minute|minutes|h|hr|hrs|hour|hours))"
|
||||
|
||||
|
||||
@ -255,26 +141,21 @@ class PluginConfigFormMixin:
|
||||
allow_crawl_execution_config_fields = True
|
||||
|
||||
def build_plugin_groups(self, runtime_config: Mapping[str, Any] | None = None) -> None:
|
||||
all_plugins = get_plugins()
|
||||
catalog = get_plugin_catalog()
|
||||
all_plugins = set(catalog)
|
||||
plugin_configs = discover_plugin_configs()
|
||||
runtime_config = runtime_config or get_config()
|
||||
self.plugin_config_binary_urls = get_plugin_config_binary_urls(runtime_config)
|
||||
grouped_plugins = set().union(*(group[-1] for group in PLUGIN_GROUP_DEFINITIONS))
|
||||
other_plugins = tuple(sorted(set(all_plugins) - grouped_plugins - HIDDEN_PLUGIN_CONFIG_UI_PLUGINS))
|
||||
grouped_plugins = catalog.groups()
|
||||
|
||||
for field_name, *_rest, plugin_names in PLUGIN_GROUP_DEFINITIONS:
|
||||
group_specs = []
|
||||
for category, field_name, title in PLUGIN_GROUPS:
|
||||
plugin_names = tuple(plugin.name for plugin in grouped_plugins.get(category, []))
|
||||
group_specs.append((field_name, title, "", "", "", plugin_names))
|
||||
if field_name in self.fields:
|
||||
get_choice_field(self, field_name).choices = [
|
||||
(p, get_plugin_choice_label(p, plugin_configs)) for p in plugin_names if p in all_plugins
|
||||
]
|
||||
|
||||
if "other_plugins" in self.fields:
|
||||
get_choice_field(self, "other_plugins").choices = [(p, get_plugin_choice_label(p, plugin_configs)) for p in other_plugins]
|
||||
|
||||
group_specs = (
|
||||
*PLUGIN_GROUP_DEFINITIONS,
|
||||
("other_plugins", "Other", "", "", "", other_plugins),
|
||||
)
|
||||
binary_url_lookup = _build_required_binary_url_lookup(plugin_configs, runtime_config)
|
||||
self.plugin_groups = [
|
||||
{
|
||||
|
||||
@ -1,63 +1,27 @@
|
||||
"""ArchiveBox adapters around the framework-free abx-dl plugin runtime.
|
||||
|
||||
Discovery and execution are owned by abx-dl. ArchiveBox keeps only the small
|
||||
Django projection adapter and its application-specific URL-output reader.
|
||||
"""
|
||||
Hook discovery and execution helpers for ArchiveBox plugins.
|
||||
|
||||
ArchiveBox no longer drives plugin execution itself during normal crawls.
|
||||
`abx-dl` owns the live runtime and emits typed bus events; ArchiveBox mainly:
|
||||
|
||||
- discovers hook files for inspection / docs / legacy direct execution helpers
|
||||
- executes individual hook scripts when explicitly requested
|
||||
- parses hook stdout JSONL records into ArchiveBox models when needed
|
||||
|
||||
Hook-backed event families are discovered from filenames like:
|
||||
on_CrawlSetup__*
|
||||
on_Snapshot__*
|
||||
|
||||
Internal bus event names are normalized to the corresponding
|
||||
`on_{EventFamily}__*` prefix by a simple string transform. If no scripts exist
|
||||
for that prefix, discovery returns `[]`.
|
||||
|
||||
Directory structure:
|
||||
abx_plugins/plugins/<plugin_name>/on_<Event>__<hook_name>.<ext> (built-in package)
|
||||
data/custom_plugins/<plugin_name>/on_<Event>__<hook_name>.<ext> (user)
|
||||
|
||||
Hook contract:
|
||||
Input: --url=<url> (and other --key=value args)
|
||||
Output: JSONL records to stdout, files to $PWD
|
||||
Exit: 0 = success, non-zero = failure
|
||||
|
||||
Execution order:
|
||||
- Hooks are named with two-digit prefixes (00-99) and sorted lexicographically by filename
|
||||
- Foreground hooks run sequentially in that order
|
||||
- Background hooks (.bg suffix) run concurrently and do not block foreground progress
|
||||
- After all foreground hooks complete, background hooks receive SIGTERM and must finalize
|
||||
|
||||
Hook naming convention:
|
||||
on_{EventFamily}__{run_order}_{description}[.bg].{ext}
|
||||
|
||||
API:
|
||||
discover_hooks(event) -> List[Path] Find hook scripts for a hook-backed event family
|
||||
run_hook(script, ...) -> Process Execute a hook script directly
|
||||
is_background_hook(name) -> bool Check if hook is background (.bg suffix)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
__package__ = "archivebox.plugins"
|
||||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Optional, Protocol, TypeGuard, runtime_checkable
|
||||
from typing import TYPE_CHECKING, Any, Protocol, TypeGuard, runtime_checkable
|
||||
|
||||
from asgiref.sync import async_to_sync
|
||||
|
||||
from abx_dl.execution import execute_hook
|
||||
from abx_dl.models import Hook, parse_hook_filename
|
||||
|
||||
from archivebox.config.constants import CONSTANTS
|
||||
from archivebox.config.version import VERSION
|
||||
from archivebox.misc.util import fix_url_from_markdown, sanitize_extracted_url
|
||||
from archivebox.plugins.discovery import (
|
||||
BUILTIN_PLUGINS_DIR,
|
||||
USER_PLUGINS_DIR,
|
||||
ConfigLookup,
|
||||
get_enabled_plugins,
|
||||
get_plugin_special_config,
|
||||
)
|
||||
from archivebox.plugins.discovery import ConfigLookup, get_enabled_plugins, get_plugin_catalog, get_plugin_special_config
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from archivebox.machine.models import Process
|
||||
@ -80,67 +44,16 @@ def _config_to_overrides(config: ConfigLookup | Mapping[str, Any] | None) -> dic
|
||||
return dict(config.items())
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Hook Step Extraction
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def is_background_hook(hook_name: str) -> bool:
|
||||
"""
|
||||
Check if a hook is a background hook (doesn't block foreground progression).
|
||||
|
||||
Background hooks have '.bg.' in their filename before the extension.
|
||||
|
||||
Args:
|
||||
hook_name: Hook filename (e.g., 'on_Snapshot__10_chrome_tab.daemon.bg.js')
|
||||
|
||||
Returns:
|
||||
True if background hook, False if foreground.
|
||||
|
||||
Examples:
|
||||
is_background_hook('on_Snapshot__10_chrome_tab.daemon.bg.js') -> True
|
||||
is_background_hook('on_Snapshot__50_wget.py') -> False
|
||||
is_background_hook('on_Snapshot__63_media.finite.bg.py') -> True
|
||||
"""
|
||||
return ".bg." in hook_name or "__background" in hook_name
|
||||
parsed = parse_hook_filename(Path(hook_name).name)
|
||||
return bool(parsed and parsed[2])
|
||||
|
||||
|
||||
def normalize_hook_event_name(event_name: str) -> str | None:
|
||||
"""
|
||||
Normalize a hook event family or event class name to its on_* prefix.
|
||||
|
||||
Examples:
|
||||
CrawlSetupEvent -> CrawlSetup
|
||||
SnapshotEvent -> Snapshot
|
||||
BinaryEvent -> Binary
|
||||
CrawlCleanupEvent -> CrawlCleanup
|
||||
"""
|
||||
normalized = str(event_name or "").strip()
|
||||
if not normalized:
|
||||
return None
|
||||
|
||||
if normalized.endswith("Event"):
|
||||
return normalized[:-5] or None
|
||||
return normalized
|
||||
|
||||
|
||||
def _model_output_dir_from_child_path(path: Path, marker: str) -> Path | None:
|
||||
"""
|
||||
Infer the model output dir from a model dir or one of its plugin subdirs.
|
||||
|
||||
Current ArchiveBox snapshot/crawl dirs are:
|
||||
.../{snapshots,crawls}/YYYYMMDD/domain/uuid[/plugin]
|
||||
"""
|
||||
parts = path.resolve().parts
|
||||
try:
|
||||
marker_index = parts.index(marker)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
model_end_index = marker_index + 4
|
||||
if len(parts) < model_end_index:
|
||||
return None
|
||||
return Path(*parts[:model_end_index])
|
||||
return normalized.removesuffix("Event") or None
|
||||
|
||||
|
||||
def discover_hooks(
|
||||
@ -149,96 +62,93 @@ def discover_hooks(
|
||||
config: ConfigLookup | None = None,
|
||||
**config_kwargs: Any,
|
||||
) -> list[Path]:
|
||||
"""
|
||||
Find all hook scripts for an event family.
|
||||
|
||||
Searches both built-in and user plugin directories.
|
||||
Filters out hooks from disabled plugins by default (respects USE_/SAVE_ flags).
|
||||
Returns scripts sorted alphabetically by filename for deterministic execution order.
|
||||
|
||||
Hook naming convention uses numeric prefixes to control order:
|
||||
on_Snapshot__10_title.py # runs first
|
||||
on_Snapshot__15_singlefile.py # runs second
|
||||
on_Snapshot__26_readability.py # runs later (depends on singlefile)
|
||||
|
||||
Args:
|
||||
event_name: Hook event family or event class name.
|
||||
Examples: 'CrawlSetupEvent', 'Snapshot'.
|
||||
Event names are normalized by stripping a trailing `Event`.
|
||||
If no matching `on_{EventFamily}__*` scripts exist, returns [].
|
||||
filter_disabled: If True, skip hooks from disabled plugins (default: True)
|
||||
config: Optional pre-merged config dict from get_config().
|
||||
**config_kwargs: Scope/override args forwarded to get_config() when config is not supplied.
|
||||
|
||||
Returns:
|
||||
Sorted list of hook script paths from enabled plugins only.
|
||||
|
||||
Examples:
|
||||
# With proper config context (recommended):
|
||||
from archivebox.config.common import get_config
|
||||
config = get_config(crawl=my_crawl, snapshot=my_snapshot)
|
||||
discover_hooks('Snapshot', config=config)
|
||||
# Returns: [Path('.../on_Snapshot__10_title.py'), ...] (wget excluded if SAVE_WGET=False)
|
||||
|
||||
# Without config (uses global defaults):
|
||||
discover_hooks('Snapshot')
|
||||
# Returns: [Path('.../on_Snapshot__10_title.py'), ...]
|
||||
|
||||
# Show all plugins regardless of enabled status:
|
||||
discover_hooks('Snapshot', filter_disabled=False)
|
||||
# Returns: [Path('.../on_Snapshot__10_title.py'), ..., Path('.../on_Snapshot__50_wget.py')]
|
||||
"""
|
||||
hook_event_name = normalize_hook_event_name(event_name)
|
||||
if not hook_event_name:
|
||||
"""Return the exact catalog hooks used by abx-dl, in execution order."""
|
||||
normalized = normalize_hook_event_name(event_name)
|
||||
if not normalized or normalized == "BinaryRequest":
|
||||
return []
|
||||
if hook_event_name == "BinaryRequest":
|
||||
return []
|
||||
|
||||
hooks = []
|
||||
|
||||
for base_dir in (BUILTIN_PLUGINS_DIR, USER_PLUGINS_DIR):
|
||||
if not base_dir.exists():
|
||||
continue
|
||||
|
||||
# Search for hook scripts in all subdirectories
|
||||
for ext in ("sh", "py", "js"):
|
||||
pattern = f"*/on_{hook_event_name}__*.{ext}"
|
||||
hooks.extend(base_dir.glob(pattern))
|
||||
|
||||
# Also check for hooks directly in the plugins directory
|
||||
pattern_direct = f"on_{hook_event_name}__*.{ext}"
|
||||
hooks.extend(base_dir.glob(pattern_direct))
|
||||
|
||||
names = None
|
||||
if filter_disabled:
|
||||
# Get merged config if not provided (lazy import to avoid circular dependency)
|
||||
if config is None:
|
||||
from archivebox.config.common import get_config
|
||||
|
||||
config = get_config(**config_kwargs)
|
||||
names = get_enabled_plugins(config=config)
|
||||
return [hook.path for _plugin, hook in get_plugin_catalog().hooks(normalized, names=names)]
|
||||
|
||||
enabled_plugins = set(get_enabled_plugins(config=config))
|
||||
enabled_hooks = []
|
||||
|
||||
for hook in hooks:
|
||||
# Get plugin name from parent directory
|
||||
# e.g., abx_plugins/plugins/wget/on_Snapshot__50_wget.py -> 'wget'
|
||||
plugin_name = hook.parent.name
|
||||
def _catalog_hook(script: Path) -> Hook:
|
||||
script = script.resolve()
|
||||
for plugin in get_plugin_catalog().values():
|
||||
for hook in plugin.hooks:
|
||||
if hook.path.resolve() == script:
|
||||
return hook
|
||||
parsed = parse_hook_filename(script.name)
|
||||
if parsed is None:
|
||||
raise ValueError(f"Not a valid plugin hook filename: {script.name}")
|
||||
event, order, is_background = parsed
|
||||
return Hook(
|
||||
name=script.name,
|
||||
event=event,
|
||||
plugin_name=script.parent.name,
|
||||
path=script,
|
||||
order=order,
|
||||
is_background=is_background,
|
||||
)
|
||||
|
||||
# Check if this is a plugin directory (not the root plugins dir)
|
||||
if hook.parent.resolve() in (BUILTIN_PLUGINS_DIR.resolve(), USER_PLUGINS_DIR.resolve()):
|
||||
# Hook is in root plugins directory, not a plugin subdir
|
||||
# Include it by default (no filtering for non-plugin hooks)
|
||||
enabled_hooks.append(hook)
|
||||
|
||||
def _hook_environment(config: ConfigLookup | Mapping[str, Any] | None, **config_scope: Any) -> tuple[dict[str, str], Any]:
|
||||
from archivebox.config.common import (
|
||||
ArchiveBoxConfig,
|
||||
_archivebox_config_input_names,
|
||||
get_config,
|
||||
normalize_runtime_config,
|
||||
)
|
||||
|
||||
overrides = _config_to_overrides(config)
|
||||
resolved = get_config(overrides=overrides, **config_scope)
|
||||
runtime = normalize_runtime_config(
|
||||
resolved.for_crawl_runtime(runtime_overrides=overrides),
|
||||
json_safe=False,
|
||||
)
|
||||
runtime.update(normalize_runtime_config(overrides, json_safe=False))
|
||||
|
||||
env = os.environ.copy()
|
||||
config_input_names = _archivebox_config_input_names()
|
||||
for key in config_input_names:
|
||||
env.pop(key, None)
|
||||
env.pop("PLUGINS", None)
|
||||
env["PATH"] = os.environ.get("PATH", "")
|
||||
env["DATA_DIR"] = str(CONSTANTS.DATA_DIR)
|
||||
env["LIBRARY_VERSION"] = VERSION
|
||||
env.setdefault("MACHINE_ID", os.environ.get("MACHINE_ID", CONSTANTS.MACHINE_ID))
|
||||
|
||||
canonical_config_keys = set(ArchiveBoxConfig.model_fields)
|
||||
for key, value in runtime.items():
|
||||
if key == "PATH" or value is None:
|
||||
continue
|
||||
if key in config_input_names and key not in canonical_config_keys:
|
||||
continue
|
||||
if isinstance(value, bool):
|
||||
env[key] = "true" if value else "false"
|
||||
elif isinstance(value, (dict, list)):
|
||||
import json
|
||||
|
||||
if plugin_name in enabled_plugins:
|
||||
enabled_hooks.append(hook)
|
||||
env[key] = json.dumps(value)
|
||||
else:
|
||||
env[key] = str(value)
|
||||
|
||||
hooks = enabled_hooks
|
||||
|
||||
# Sort by filename (not full path) to ensure numeric prefix ordering works
|
||||
# e.g., on_Snapshot__10_title.py sorts before on_Snapshot__26_readability.py
|
||||
return sorted(set(hooks), key=lambda p: p.name)
|
||||
node_modules_dir = runtime.get("NODE_MODULES_DIR")
|
||||
lib_dir = runtime.get("ABXPKG_LIB_DIR")
|
||||
if not node_modules_dir and lib_dir:
|
||||
node_modules_dir = Path(lib_dir) / "pnpm" / "packages" / "chrome" / "node_modules"
|
||||
if node_modules_dir:
|
||||
env["NODE_MODULES_DIR"] = str(node_modules_dir)
|
||||
env["NODE_MODULE_DIR"] = str(node_modules_dir)
|
||||
node_path = [part for part in str(runtime.get("NODE_PATH") or "").split(os.pathsep) if part]
|
||||
if str(node_modules_dir) not in node_path:
|
||||
node_path.append(str(node_modules_dir))
|
||||
env["NODE_PATH"] = os.pathsep.join(node_path)
|
||||
return env, resolved
|
||||
|
||||
|
||||
def run_hook(
|
||||
@ -246,384 +156,83 @@ def run_hook(
|
||||
output_dir: Path,
|
||||
config: ConfigLookup | Mapping[str, Any] | None = None,
|
||||
timeout: int | None = None,
|
||||
parent: Optional["Process"] = None,
|
||||
parent: Process | None = None,
|
||||
**kwargs: Any,
|
||||
) -> "Process":
|
||||
"""
|
||||
Execute a hook script with the given arguments using Process model.
|
||||
|
||||
This is the low-level hook executor that creates a Process record and
|
||||
uses Process.launch() for subprocess management.
|
||||
|
||||
Config is passed to hooks via environment variables. Crawl/snapshot callers
|
||||
should pass the runtime config produced by for_crawl_runtime().
|
||||
|
||||
Args:
|
||||
script: Path to the hook script (.sh, .py, or .js)
|
||||
output_dir: Working directory for the script (where output files go)
|
||||
config: Optional runtime config dict from for_crawl_runtime().
|
||||
If omitted, pass scope/override args using kwargs prefixed with config_.
|
||||
timeout: Maximum execution time in seconds
|
||||
If None, auto-detects from PLUGINNAME_TIMEOUT config (fallback to TIMEOUT, default 300)
|
||||
parent: Optional parent Process (for tracking worker->hook hierarchy)
|
||||
**kwargs: Arguments passed to the script as --key=value
|
||||
|
||||
Returns:
|
||||
Process model instance (use process.exit_code, process.stdout, process.get_records())
|
||||
|
||||
Example:
|
||||
from archivebox.config.common import get_config
|
||||
config = get_config(crawl=my_crawl, snapshot=my_snapshot).for_crawl_runtime(crawl=my_crawl, snapshot=my_snapshot)
|
||||
process = run_hook(hook_path, output_dir, config=config, url=url, snapshot_id=id)
|
||||
if process.status == 'exited':
|
||||
records = process.get_records() # Get parsed JSONL output
|
||||
"""
|
||||
from archivebox.machine.models import Process, Machine, NetworkInterface
|
||||
from archivebox.config.common import (
|
||||
ArchiveBoxConfig,
|
||||
_archivebox_config_input_names,
|
||||
get_config,
|
||||
normalize_runtime_config,
|
||||
_plugin_enabled_config_keys,
|
||||
)
|
||||
) -> Process:
|
||||
"""Compatibility adapter for finite direct calls; abx-dl owns execution."""
|
||||
from archivebox.machine.models import Process
|
||||
from archivebox.services.process_service import ProcessService as PersistedProcessService
|
||||
from archivebox.services.process_service import parse_event_datetime
|
||||
from abx_dl.orchestrator import create_bus
|
||||
|
||||
if parent is not None:
|
||||
kwargs.setdefault("_parent_process_id", str(parent.id))
|
||||
config_scope = {key.removeprefix("config_"): kwargs.pop(key) for key in list(kwargs) if key.startswith("config_")}
|
||||
config_overrides = _config_to_overrides(config)
|
||||
explicit_override_keys = set(config_overrides)
|
||||
resolved_config = get_config(overrides=config_overrides, **config_scope)
|
||||
hook_config = normalize_runtime_config(
|
||||
resolved_config.for_crawl_runtime(runtime_overrides=config_overrides),
|
||||
json_safe=False,
|
||||
)
|
||||
hook_config.update(normalize_runtime_config(config_overrides, json_safe=False))
|
||||
plugin_enabled_keys = set(_plugin_enabled_config_keys().values())
|
||||
if plugin_enabled_keys.intersection(hook_config):
|
||||
for enabled_key in plugin_enabled_keys:
|
||||
hook_config.setdefault(enabled_key, False)
|
||||
|
||||
# Auto-detect timeout from plugin config if not explicitly provided
|
||||
env, resolved = _hook_environment(config, **config_scope)
|
||||
hook = _catalog_hook(script)
|
||||
if timeout is None:
|
||||
plugin_name = script.parent.name
|
||||
plugin_config = get_plugin_special_config(plugin_name, resolved_config)
|
||||
timeout = plugin_config["timeout"]
|
||||
if timeout:
|
||||
timeout = min(int(timeout), int(CONSTANTS.MAX_HOOK_RUNTIME_SECONDS))
|
||||
timeout = get_plugin_special_config(hook.plugin_name, resolved)["timeout"]
|
||||
timeout = min(int(timeout or 300), int(CONSTANTS.MAX_HOOK_RUNTIME_SECONDS))
|
||||
|
||||
# Get current machine
|
||||
machine = Machine.current()
|
||||
iface = NetworkInterface.current(refresh=True)
|
||||
machine = iface.machine
|
||||
bus = create_bus(name=f"ArchiveBoxHook_{hook.plugin_name}", total_timeout=float(timeout) + 30.0)
|
||||
PersistedProcessService(bus)
|
||||
|
||||
# Auto-detect parent process if not explicitly provided
|
||||
# This enables automatic hierarchy tracking: Worker -> Hook
|
||||
if parent is None:
|
||||
async def execute_and_close():
|
||||
try:
|
||||
parent = Process.current()
|
||||
except Exception:
|
||||
# If Process.current() fails (e.g., not in a worker context), leave parent as None
|
||||
pass
|
||||
|
||||
if not script.is_file():
|
||||
raise FileNotFoundError(f"Hook script not found: {script}")
|
||||
|
||||
# Hooks are opaque executables. Their shipped abxpkg shebang owns runtime
|
||||
# and dependency resolution just as it does under abx-dl.
|
||||
cmd = [str(script)]
|
||||
|
||||
# Build CLI arguments from kwargs
|
||||
for key, value in kwargs.items():
|
||||
# Skip keys that start with underscore (internal parameters)
|
||||
if key.startswith("_"):
|
||||
continue
|
||||
|
||||
arg_key = f"--{key.replace('_', '-')}"
|
||||
if isinstance(value, bool):
|
||||
if value:
|
||||
cmd.append(arg_key)
|
||||
elif value is not None and value != "":
|
||||
# JSON-encode complex values, use str for simple ones
|
||||
# Skip empty strings to avoid --key= which breaks argument parsers
|
||||
if isinstance(value, (dict, list)):
|
||||
cmd.append(f"{arg_key}={json.dumps(value)}")
|
||||
else:
|
||||
# Ensure value is converted to string and strip whitespace
|
||||
str_value = str(value).strip()
|
||||
if str_value: # Only add if non-empty after stripping
|
||||
cmd.append(f"{arg_key}={str_value}")
|
||||
|
||||
# Set up environment with base paths
|
||||
env = os.environ.copy()
|
||||
archivebox_config_input_names = _archivebox_config_input_names()
|
||||
for key in archivebox_config_input_names:
|
||||
env.pop(key, None)
|
||||
env.pop("PLUGINS", None)
|
||||
env["DATA_DIR"] = str(CONSTANTS.DATA_DIR)
|
||||
env["LIBRARY_VERSION"] = VERSION
|
||||
env.setdefault("MACHINE_ID", os.environ.get("MACHINE_ID", CONSTANTS.MACHINE_ID))
|
||||
snap_dir = hook_config.get("SNAP_DIR") or _model_output_dir_from_child_path(output_dir, CONSTANTS.SNAPSHOTS_DIR_NAME)
|
||||
crawl_dir = hook_config.get("CRAWL_DIR") or _model_output_dir_from_child_path(output_dir, CONSTANTS.CRAWLS_DIR_NAME)
|
||||
if snap_dir:
|
||||
env["SNAP_DIR"] = str(snap_dir)
|
||||
if crawl_dir:
|
||||
env["CRAWL_DIR"] = str(crawl_dir)
|
||||
|
||||
# Export the runtime library root; abx-dl/abxpkg own executable lookup env.
|
||||
lib_dir = hook_config.get("ABXPKG_LIB_DIR")
|
||||
if lib_dir:
|
||||
env["ABXPKG_LIB_DIR"] = str(lib_dir)
|
||||
|
||||
# Set Node.js module resolution paths.
|
||||
# NODE_PATH may be a path list, but NODE_MODULES_DIR is a single canonical directory.
|
||||
node_modules_dir = hook_config.get("NODE_MODULES_DIR")
|
||||
if lib_dir and "ABXPKG_LIB_DIR" in explicit_override_keys and "NODE_MODULES_DIR" not in explicit_override_keys:
|
||||
node_modules_dir = Path(lib_dir) / "pnpm" / "packages" / "chrome" / "node_modules"
|
||||
elif not node_modules_dir and lib_dir:
|
||||
node_modules_dir = Path(lib_dir) / "pnpm" / "packages" / "chrome" / "node_modules"
|
||||
|
||||
node_path_parts = [part for part in str(hook_config.get("NODE_PATH") or "").split(os.pathsep) if part]
|
||||
if node_modules_dir:
|
||||
node_modules_dir = Path(node_modules_dir)
|
||||
node_modules_dir.mkdir(parents=True, exist_ok=True)
|
||||
node_modules_dir_str = str(node_modules_dir)
|
||||
env["NODE_MODULES_DIR"] = node_modules_dir_str
|
||||
env["NODE_MODULE_DIR"] = node_modules_dir_str
|
||||
if node_modules_dir_str not in node_path_parts:
|
||||
node_path_parts.append(node_modules_dir_str)
|
||||
if node_path_parts:
|
||||
env["NODE_PATH"] = os.pathsep.join(node_path_parts)
|
||||
|
||||
# Export all config values to environment (already merged by get_config())
|
||||
# Skip keys we've already handled specially above (PATH, ABXPKG_LIB_DIR, NODE_PATH, etc.)
|
||||
SKIP_KEYS = {
|
||||
"PATH",
|
||||
"ABXPKG_LIB_DIR",
|
||||
"NODE_PATH",
|
||||
"NODE_MODULES_DIR",
|
||||
"NODE_MODULE_DIR",
|
||||
"DATA_DIR",
|
||||
"MACHINE_ID",
|
||||
"SNAP_DIR",
|
||||
"CRAWL_DIR",
|
||||
}
|
||||
canonical_config_keys = set(ArchiveBoxConfig.model_fields)
|
||||
for key, value in hook_config.items():
|
||||
if key in SKIP_KEYS:
|
||||
continue # Already handled specially above, don't overwrite
|
||||
if key in archivebox_config_input_names and key not in canonical_config_keys:
|
||||
continue
|
||||
if value is None:
|
||||
continue
|
||||
elif isinstance(value, bool):
|
||||
env[key] = "true" if value else "false"
|
||||
elif isinstance(value, (list, dict)):
|
||||
env[key] = json.dumps(value)
|
||||
else:
|
||||
env[key] = str(value)
|
||||
|
||||
# Create output directory if needed
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Detect if this is a background hook.
|
||||
# Background hooks use the .bg. filename marker.
|
||||
# Old convention: __background in stem (for backwards compatibility)
|
||||
is_background = ".bg." in script.name or "__background" in script.stem
|
||||
|
||||
try:
|
||||
# Create Process record
|
||||
process = Process.objects.create(
|
||||
machine=machine,
|
||||
iface=iface,
|
||||
parent=parent,
|
||||
process_type=Process.TypeChoices.HOOK,
|
||||
pwd=str(output_dir),
|
||||
cmd=cmd,
|
||||
return await execute_hook(
|
||||
hook,
|
||||
output_dir=output_dir,
|
||||
env=env,
|
||||
arguments=kwargs,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
# Copy the env dict we already built (includes os.environ + all customizations)
|
||||
process.env = env.copy()
|
||||
process.hydrate_binary_from_context(plugin_name=script.parent.name, hook_path=str(script))
|
||||
|
||||
# Save env before launching
|
||||
process.save()
|
||||
|
||||
# Launch subprocess using Process.launch()
|
||||
process.launch(background=is_background)
|
||||
|
||||
# Return Process object (caller can use process.exit_code, process.stdout, process.get_records())
|
||||
return process
|
||||
|
||||
except Exception as e:
|
||||
# Create a failed Process record for exceptions
|
||||
process = Process.objects.create(
|
||||
machine=machine,
|
||||
iface=iface,
|
||||
bus=bus,
|
||||
process_type=Process.TypeChoices.HOOK,
|
||||
pwd=str(output_dir),
|
||||
cmd=cmd,
|
||||
timeout=timeout,
|
||||
status=Process.StatusChoices.EXITED,
|
||||
exit_code=1,
|
||||
stderr=f"Failed to run hook: {type(e).__name__}: {e}",
|
||||
)
|
||||
finally:
|
||||
await bus.wait_until_idle()
|
||||
await bus.destroy(clear=False)
|
||||
|
||||
completed = async_to_sync(execute_and_close)()
|
||||
started_at = parse_event_datetime(completed.start_ts)
|
||||
process = Process.objects.filter(pid=completed.pid or None, started_at=started_at).order_by("-modified_at").first()
|
||||
if process is None:
|
||||
raise RuntimeError(f"Hook {hook.full_name} completed without an ArchiveBox Process projection")
|
||||
return process
|
||||
|
||||
|
||||
def extract_records_from_process(process: "Process") -> list[dict[str, Any]]:
|
||||
"""
|
||||
Extract JSONL records from a Process's stdout.
|
||||
|
||||
Adds plugin metadata to each record.
|
||||
|
||||
Args:
|
||||
process: Process model instance with stdout captured
|
||||
|
||||
Returns:
|
||||
List of parsed JSONL records with plugin metadata
|
||||
"""
|
||||
def extract_records_from_process(process: Process) -> list[dict[str, Any]]:
|
||||
"""Return hook JSONL records with generic catalog identity attached."""
|
||||
records = process.get_records()
|
||||
if not records:
|
||||
return []
|
||||
|
||||
# Extract plugin metadata from process.pwd and the shipped hook path in cmd.
|
||||
# Python hooks execute directly through their shebang, while JS and shell
|
||||
# hooks have an interpreter in cmd[0].
|
||||
plugin_name = Path(process.pwd).name if process.pwd else "unknown"
|
||||
plugin_hook = next((str(arg) for arg in process.cmd if Path(str(arg)).name.startswith("on_")), "")
|
||||
hook_name = Path(plugin_hook).name if plugin_hook else "unknown"
|
||||
|
||||
for record in records:
|
||||
# Add plugin metadata to record
|
||||
record.setdefault("plugin", plugin_name)
|
||||
record.setdefault("hook_name", hook_name)
|
||||
record.setdefault("plugin_hook", plugin_hook)
|
||||
|
||||
return records
|
||||
|
||||
|
||||
def collect_urls_from_plugins(snapshot_dir: Path) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Collect all urls.jsonl entries from parser plugin output subdirectories.
|
||||
|
||||
Each parser plugin outputs urls.jsonl to its own subdir:
|
||||
snapshot_dir/parse_rss_urls/urls.jsonl
|
||||
snapshot_dir/parse_html_urls/urls.jsonl
|
||||
etc.
|
||||
|
||||
This is not special handling - urls.jsonl is just a normal output file.
|
||||
This utility collects them all for the crawl system.
|
||||
"""
|
||||
urls = []
|
||||
|
||||
# Look in each immediate subdirectory for urls.jsonl
|
||||
"""Read the durable urls.jsonl interface emitted by parser plugins."""
|
||||
urls: list[dict[str, Any]] = []
|
||||
if not snapshot_dir.exists():
|
||||
return urls
|
||||
|
||||
for subdir in snapshot_dir.iterdir():
|
||||
if not subdir.is_dir():
|
||||
continue
|
||||
|
||||
urls_file = subdir / "urls.jsonl"
|
||||
if not urls_file.exists():
|
||||
continue
|
||||
|
||||
try:
|
||||
from archivebox.machine.models import Process
|
||||
|
||||
text = urls_file.read_text()
|
||||
for entry in Process.parse_records_from_text(text):
|
||||
if entry.get("url"):
|
||||
entry["url"] = sanitize_extracted_url(fix_url_from_markdown(str(entry["url"]).strip()))
|
||||
if not entry["url"]:
|
||||
for subdir in snapshot_dir.iterdir():
|
||||
urls_file = subdir / "urls.jsonl"
|
||||
if not subdir.is_dir() or not urls_file.is_file():
|
||||
continue
|
||||
# Track which parser plugin found this URL
|
||||
try:
|
||||
for entry in Process.parse_records_from_text(urls_file.read_text()):
|
||||
if not entry.get("url"):
|
||||
continue
|
||||
entry["url"] = sanitize_extracted_url(fix_url_from_markdown(str(entry["url"]).strip()))
|
||||
if entry["url"]:
|
||||
entry["plugin"] = subdir.name
|
||||
urls.append(entry)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
except (OSError, UnicodeError):
|
||||
continue
|
||||
return urls
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Hook Result Processing Helpers
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def process_hook_records(records: list[dict[str, Any]], overrides: dict[str, Any] | None = None) -> dict[str, int]:
|
||||
"""
|
||||
Process JSONL records emitted by hook stdout.
|
||||
|
||||
This handles hook-emitted record types such as Snapshot, Tag, and Binary.
|
||||
It does not process internal bus lifecycle events, since those
|
||||
are not emitted as JSONL records by hook subprocesses.
|
||||
|
||||
Args:
|
||||
records: List of JSONL record dicts from result['records']
|
||||
overrides: Dict with 'snapshot', 'crawl', 'dependency', 'created_by_id', etc.
|
||||
|
||||
Returns:
|
||||
Dict with counts by record type
|
||||
"""
|
||||
stats = {}
|
||||
overrides = overrides or {}
|
||||
|
||||
for record in records:
|
||||
record_type = record.get("type")
|
||||
if not record_type:
|
||||
continue
|
||||
|
||||
# Skip ArchiveResult records (they update the calling ArchiveResult, not create new ones)
|
||||
if record_type == "ArchiveResult":
|
||||
continue
|
||||
|
||||
try:
|
||||
# Dispatch to appropriate model's from_json() method
|
||||
if record_type == "Snapshot":
|
||||
from archivebox.core.models import Snapshot
|
||||
|
||||
if record.get("url"):
|
||||
record = {
|
||||
**record,
|
||||
"url": sanitize_extracted_url(fix_url_from_markdown(str(record["url"]).strip())),
|
||||
}
|
||||
if not record["url"]:
|
||||
continue
|
||||
|
||||
# Check if discovered snapshot exceeds crawl max_depth
|
||||
snapshot_depth = record.get("depth", 0)
|
||||
crawl = overrides.get("crawl")
|
||||
if crawl and snapshot_depth > crawl.max_depth:
|
||||
# Skip - this URL was discovered but exceeds max crawl depth
|
||||
continue
|
||||
|
||||
obj = Snapshot.from_json(record.copy(), overrides)
|
||||
if obj:
|
||||
stats["Snapshot"] = stats.get("Snapshot", 0) + 1
|
||||
|
||||
elif record_type == "Tag":
|
||||
from archivebox.core.models import Tag
|
||||
|
||||
obj = Tag.from_json(record.copy(), overrides)
|
||||
if obj:
|
||||
stats["Tag"] = stats.get("Tag", 0) + 1
|
||||
|
||||
elif record_type == "Binary":
|
||||
from archivebox.machine.models import Binary
|
||||
|
||||
obj = Binary.from_json(record.copy(), overrides)
|
||||
if obj:
|
||||
stats[record_type] = stats.get(record_type, 0) + 1
|
||||
|
||||
else:
|
||||
import sys
|
||||
|
||||
print(f"Warning: Unknown record type '{record_type}' from hook output", file=sys.stderr)
|
||||
|
||||
except Exception as e:
|
||||
import sys
|
||||
|
||||
print(f"Warning: Failed to create {record_type}: {e}", file=sys.stderr)
|
||||
continue
|
||||
|
||||
return stats
|
||||
|
||||
@ -8,19 +8,17 @@ import re
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterable
|
||||
from contextlib import contextmanager
|
||||
from functools import wraps
|
||||
from pathlib import Path
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
from typing import Any
|
||||
|
||||
from asgiref.sync import sync_to_async
|
||||
from django.db import IntegrityError
|
||||
from django.utils import timezone
|
||||
|
||||
from abx_dl.events import PROCESS_EXIT_SKIPPED, ArchiveResultEvent, ProcessCompletedEvent, ProcessStartedEvent, SnapshotEvent
|
||||
from abx_dl.output_files import guess_mimetype
|
||||
from abx_dl.output_files import OutputManifest
|
||||
from abx_dl.services.base import BaseService
|
||||
|
||||
from .process_service import parse_event_datetime
|
||||
@ -72,129 +70,15 @@ def _perf_span(label: str):
|
||||
print(f"PERF_TRACE label={label} ms={elapsed_ms:.3f}", file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ModelDumpable(Protocol):
|
||||
def model_dump(self) -> dict[str, Any]: ...
|
||||
|
||||
|
||||
def _collect_output_metadata(plugin_dir: Path) -> tuple[dict[str, dict], int, str]:
|
||||
exclude_names = {"stdout.log", "stderr.log", "process.pid", "hook.pid", "listener.pid"}
|
||||
output_files: dict[str, dict] = {}
|
||||
mime_sizes: dict[str, int] = defaultdict(int)
|
||||
total_size = 0
|
||||
|
||||
if not plugin_dir.exists():
|
||||
return output_files, total_size, ""
|
||||
|
||||
for file_path in plugin_dir.rglob("*"):
|
||||
if not file_path.is_file():
|
||||
continue
|
||||
if ".hooks" in file_path.parts:
|
||||
continue
|
||||
if file_path.name in exclude_names:
|
||||
continue
|
||||
try:
|
||||
stat = file_path.stat()
|
||||
except OSError:
|
||||
continue
|
||||
mime_type = guess_mimetype(file_path) or "application/octet-stream"
|
||||
relative_path = str(file_path.relative_to(plugin_dir))
|
||||
output_files[relative_path] = {
|
||||
"extension": file_path.suffix.lower().lstrip("."),
|
||||
"mimetype": mime_type,
|
||||
"size": stat.st_size,
|
||||
}
|
||||
mime_sizes[mime_type] += stat.st_size
|
||||
total_size += stat.st_size
|
||||
|
||||
output_mimetypes = ",".join(mime for mime, _size in sorted(mime_sizes.items(), key=lambda item: item[1], reverse=True))
|
||||
return output_files, total_size, output_mimetypes
|
||||
|
||||
|
||||
def _coerce_output_file_size(value: Any) -> int:
|
||||
try:
|
||||
return max(int(value or 0), 0)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def _normalize_output_files(raw_output_files: Any) -> dict[str, dict]:
|
||||
def _enrich_metadata(path: str, metadata: dict[str, Any]) -> dict[str, Any]:
|
||||
normalized = dict(metadata)
|
||||
if "extension" not in normalized:
|
||||
normalized["extension"] = Path(path).suffix.lower().lstrip(".")
|
||||
if "mimetype" not in normalized:
|
||||
guessed = guess_mimetype(path)
|
||||
if guessed:
|
||||
normalized["mimetype"] = guessed
|
||||
return normalized
|
||||
|
||||
if raw_output_files is None:
|
||||
return {}
|
||||
|
||||
if isinstance(raw_output_files, str):
|
||||
try:
|
||||
raw_output_files = json.loads(raw_output_files)
|
||||
except json.JSONDecodeError:
|
||||
return {}
|
||||
|
||||
if isinstance(raw_output_files, dict):
|
||||
normalized: dict[str, dict] = {}
|
||||
for path, metadata in raw_output_files.items():
|
||||
if not path:
|
||||
continue
|
||||
metadata_dict = dict(metadata) if isinstance(metadata, dict) else {}
|
||||
metadata_dict.pop("path", None)
|
||||
normalized[str(path)] = _enrich_metadata(str(path), metadata_dict)
|
||||
return normalized
|
||||
|
||||
if not isinstance(raw_output_files, Iterable):
|
||||
return {}
|
||||
|
||||
normalized: dict[str, dict] = {}
|
||||
for item in raw_output_files:
|
||||
if isinstance(item, str):
|
||||
normalized[item] = _enrich_metadata(item, {})
|
||||
continue
|
||||
if isinstance(item, ModelDumpable):
|
||||
item = item.model_dump()
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
path = str(item.get("path") or "").strip()
|
||||
if not path:
|
||||
continue
|
||||
normalized[path] = _enrich_metadata(path, {key: value for key, value in item.items() if key != "path" and value not in (None, "")})
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
def _has_structured_output_metadata(output_files: dict[str, dict]) -> bool:
|
||||
return any(any(key in metadata for key in ("extension", "mimetype", "size")) for metadata in output_files.values())
|
||||
|
||||
|
||||
def _summarize_output_files(output_files: dict[str, dict]) -> tuple[int, str]:
|
||||
mime_sizes: dict[str, int] = defaultdict(int)
|
||||
total_size = 0
|
||||
|
||||
for metadata in output_files.values():
|
||||
if not isinstance(metadata, dict):
|
||||
continue
|
||||
size = _coerce_output_file_size(metadata.get("size"))
|
||||
mimetype = str(metadata.get("mimetype") or "").strip()
|
||||
total_size += size
|
||||
if mimetype and size:
|
||||
mime_sizes[mimetype] += size
|
||||
|
||||
output_mimetypes = ",".join(mime for mime, _size in sorted(mime_sizes.items(), key=lambda item: item[1], reverse=True))
|
||||
return total_size, output_mimetypes
|
||||
def _manifest_metadata(manifest: OutputManifest) -> tuple[dict[str, dict], int, str]:
|
||||
return manifest.as_mapping(), manifest.total_size, ",".join(manifest.mimetypes)
|
||||
|
||||
|
||||
def _resolve_output_metadata(raw_output_files: Any, plugin_dir: Path) -> tuple[dict[str, dict], int, str]:
|
||||
normalized_output_files = _normalize_output_files(raw_output_files)
|
||||
if normalized_output_files and _has_structured_output_metadata(normalized_output_files):
|
||||
output_size, output_mimetypes = _summarize_output_files(normalized_output_files)
|
||||
return normalized_output_files, output_size, output_mimetypes
|
||||
return _collect_output_metadata(plugin_dir)
|
||||
manifest = OutputManifest.from_value(raw_output_files)
|
||||
if manifest.files and any(output_file.size for output_file in manifest.files):
|
||||
return _manifest_metadata(manifest)
|
||||
return _manifest_metadata(OutputManifest.scan(plugin_dir, containment_root=plugin_dir.parent))
|
||||
|
||||
|
||||
def _normalize_status(status: str) -> str:
|
||||
@ -362,8 +246,9 @@ def _save_archiveresult_event_to_db(
|
||||
|
||||
if result.output_files:
|
||||
merged_output_files = {**result.output_files, **defaults["output_files"]}
|
||||
defaults["output_files"] = merged_output_files
|
||||
defaults["output_size"], defaults["output_mimetypes"] = _summarize_output_files(merged_output_files)
|
||||
defaults["output_files"], defaults["output_size"], defaults["output_mimetypes"] = _manifest_metadata(
|
||||
OutputManifest.from_value(merged_output_files),
|
||||
)
|
||||
defaults["output_size"] = max(defaults["output_size"], int(result.output_size or 0))
|
||||
defaults["output_mimetypes"] = ",".join(
|
||||
dict.fromkeys(
|
||||
|
||||
@ -21,7 +21,6 @@ from rich.console import Console
|
||||
from rich.text import Text
|
||||
|
||||
from abxpkg.binary_service import BinaryRequestEvent, BinaryService
|
||||
from abx_dl.config import GlobalConfig, RuntimeConfig
|
||||
from abx_dl.events import (
|
||||
CrawlAbortEvent,
|
||||
CrawlCleanupEvent,
|
||||
@ -38,9 +37,10 @@ from abx_dl.events import (
|
||||
)
|
||||
from abx_dl.heartbeat import CrawlHeartbeat
|
||||
from abx_dl.limits import CrawlLimitState
|
||||
from abx_dl.models import Plugin, Snapshot as AbxSnapshot, discover_plugins, filter_plugins
|
||||
from abx_dl.catalog import PluginCatalog
|
||||
from abx_dl.models import Plugin, Snapshot as AbxSnapshot, filter_plugins
|
||||
from abx_dl.orchestrator import (
|
||||
compute_phase_timeout,
|
||||
ExecutionPlan,
|
||||
create_bus,
|
||||
install_plugins as abx_install_plugins,
|
||||
setup_services as setup_abx_services,
|
||||
@ -59,6 +59,7 @@ from archivebox.config.common import (
|
||||
)
|
||||
from archivebox.misc.db import run_db_analyze_batch
|
||||
from archivebox.core.shutdown_util import foreground_shutdown_signals, raise_if_shutdown_requested
|
||||
from archivebox.plugins.discovery import get_plugin_catalog
|
||||
from archivebox.search.sonic_daemon import register_sonic_daemon_event_handler
|
||||
from archivebox.workers.models import ACTIVE_STATE_LEASE_SECONDS
|
||||
from archivebox.crawls.locks import crawl_lifecycle_lock
|
||||
@ -128,7 +129,11 @@ def _is_nonfatal_setup_hook(plugin_name: str, hook_name: str) -> bool:
|
||||
|
||||
|
||||
def _discover_archivebox_plugins() -> dict[str, Plugin]:
|
||||
return discover_plugins(runtime="archivebox")
|
||||
return _discover_archivebox_catalog().plugins
|
||||
|
||||
|
||||
def _discover_archivebox_catalog() -> PluginCatalog:
|
||||
return get_plugin_catalog()
|
||||
|
||||
|
||||
def _runner_task_context() -> contextvars.Context:
|
||||
@ -217,7 +222,8 @@ class CrawlRunner:
|
||||
):
|
||||
self.crawl = crawl
|
||||
self.bus = create_bus(name=_bus_name("ArchiveBox", str(crawl.id)), total_timeout=3600.0)
|
||||
self.plugins = _discover_archivebox_plugins()
|
||||
self.catalog = _discover_archivebox_catalog()
|
||||
self.plugins = self.catalog.plugins
|
||||
HookProcessService(self.bus, emit_jsonl=False, interactive_tty=interactive_interrupts)
|
||||
register_sonic_daemon_event_handler(self.bus)
|
||||
PersistedProcessService(self.bus)
|
||||
@ -307,7 +313,7 @@ class CrawlRunner:
|
||||
return
|
||||
|
||||
def runtime_plugins(self) -> dict[str, Plugin]:
|
||||
return filter_plugins(self.plugins, self.selected_plugins, include_providers=True) if self.selected_plugins else self.plugins
|
||||
return self.catalog.select(self.selected_plugins).plugins if self.selected_plugins else self.plugins
|
||||
|
||||
@property
|
||||
def allow_maintenance_on_inactive_crawl(self) -> bool:
|
||||
@ -359,11 +365,6 @@ class CrawlRunner:
|
||||
with live_ui if live_ui is not None else nullcontext():
|
||||
try:
|
||||
await heartbeat.start()
|
||||
await _emit_machine_config(
|
||||
self.bus,
|
||||
config=self.base_config,
|
||||
derived_config=self.derived_config,
|
||||
)
|
||||
if snapshot_ids:
|
||||
root_snapshot_id = snapshot_ids[0]
|
||||
await self.run_crawl(root_snapshot_id, snapshot_ids)
|
||||
@ -562,7 +563,7 @@ class CrawlRunner:
|
||||
def load_run_state(self) -> list[str]:
|
||||
from archivebox.config.common import get_config
|
||||
from archivebox.core.models import Snapshot
|
||||
from archivebox.plugins.hooks import discover_hooks
|
||||
from archivebox.plugins.discovery import get_enabled_plugins
|
||||
from archivebox.machine.models import Machine, NetworkInterface, Process
|
||||
|
||||
self.primary_url = self.crawl.get_urls_list()[0] if self.crawl.get_urls_list() else ""
|
||||
@ -588,9 +589,10 @@ class CrawlRunner:
|
||||
if raw_plugins:
|
||||
self.selected_plugins = [name.strip() for name in raw_plugins.split(",") if name.strip()]
|
||||
else:
|
||||
enabled_plugins = get_enabled_plugins(config=self.base_config)
|
||||
runtime_events = ("CrawlSetup", "CrawlCleanup", "Snapshot", "SnapshotCleanup")
|
||||
runtime_plugins = {
|
||||
hook.parent.name for event_name in runtime_events for hook in discover_hooks(event_name, config=self.base_config)
|
||||
plugin.name for event_name in runtime_events for plugin, _hook in self.catalog.hooks(event_name, names=enabled_plugins)
|
||||
}
|
||||
self.selected_plugins = sorted(runtime_plugins) or None
|
||||
if self.crawl.is_paused:
|
||||
@ -690,7 +692,7 @@ class CrawlRunner:
|
||||
if crawl.is_finished():
|
||||
if crawl.status != Crawl.StatusChoices.SEALED:
|
||||
if crawl.status == Crawl.StatusChoices.STARTED:
|
||||
crawl.sm.seal()
|
||||
crawl.seal()
|
||||
else:
|
||||
crawl.update_and_requeue(
|
||||
status=Crawl.StatusChoices.SEALED,
|
||||
@ -834,18 +836,23 @@ class CrawlRunner:
|
||||
config = normalize_runtime_config(snapshot["config"])
|
||||
derived_config = normalize_runtime_config(self.derived_config)
|
||||
output_dir = Path(self.crawl_output_dir)
|
||||
plugins = self.runtime_plugins()
|
||||
plan = ExecutionPlan.build(
|
||||
self.catalog,
|
||||
selected_plugins=self.selected_plugins,
|
||||
config=config,
|
||||
derived_config=derived_config,
|
||||
runtime="archivebox",
|
||||
)
|
||||
setup_hooks = [(plugin, hook) for plugin in plan.plugins.values() for hook in plugin.filter_hooks("CrawlSetup")]
|
||||
abx_snapshot = AbxSnapshot(
|
||||
id=snapshot["id"],
|
||||
url=snapshot["url"],
|
||||
depth=int(snapshot["depth"]),
|
||||
crawl_id=str(self.crawl.id),
|
||||
)
|
||||
setup_hooks = [(plugin, hook) for plugin in plugins.values() for hook in plugin.filter_hooks("CrawlSetup")]
|
||||
crawl_setup_phase_timeout = compute_phase_timeout(setup_hooks, config)
|
||||
snapshot_hooks = [(plugin, hook) for plugin in plugins.values() for hook in plugin.filter_hooks("Snapshot")]
|
||||
crawl_setup_phase_timeout = plan.crawl_setup_timeout
|
||||
max_snapshot_count = max(1, int(config.get("CRAWL_MAX_URLS") or len(snapshot_ids) or 1))
|
||||
snapshot_phase_timeout = compute_phase_timeout(snapshot_hooks, config) + 120.0
|
||||
snapshot_phase_timeout = plan.snapshot_timeout + 120.0
|
||||
all_snapshots_phase_timeout = snapshot_phase_timeout * max_snapshot_count
|
||||
crawl_cleanup_phase_timeout = crawl_setup_phase_timeout
|
||||
crawl_lifecycle_timeout = (
|
||||
@ -855,10 +862,9 @@ class CrawlRunner:
|
||||
+ CrawlCompletedEvent.model_fields["event_timeout"].default
|
||||
+ 30.0
|
||||
)
|
||||
await _emit_machine_config(self.bus, config=config, derived_config=derived_config)
|
||||
setup_abx_services(
|
||||
await plan.seed_config(self.bus)
|
||||
plan.attach_services(
|
||||
self.bus,
|
||||
plugins=plugins,
|
||||
url=snapshot["url"],
|
||||
snapshot=abx_snapshot,
|
||||
output_dir=output_dir,
|
||||
@ -869,10 +875,6 @@ class CrawlRunner:
|
||||
snapshot_cleanup_enabled=False,
|
||||
crawl_cleanup_enabled=True,
|
||||
crawl_completed_enabled=False,
|
||||
crawl_setup_phase_timeout=crawl_setup_phase_timeout,
|
||||
snapshot_phase_timeout=crawl_setup_phase_timeout,
|
||||
snapshot_cleanup_phase_timeout=crawl_setup_phase_timeout,
|
||||
crawl_cleanup_phase_timeout=crawl_setup_phase_timeout,
|
||||
auto_install=True,
|
||||
emit_jsonl=False,
|
||||
abort_requested=self.crawl_is_cancelled,
|
||||
@ -1131,24 +1133,26 @@ class CrawlRunner:
|
||||
depth=int(snapshot["depth"]),
|
||||
crawl_id=str(self.crawl.id),
|
||||
)
|
||||
snapshot_hooks = [(plugin, hook) for plugin in plugins.values() for hook in plugin.filter_hooks("Snapshot")]
|
||||
snapshot_phase_timeout = compute_phase_timeout(snapshot_hooks, config) + 120.0
|
||||
await _emit_machine_config(self.bus, config=config, derived_config=derived_config, parent_event=crawl_start_event)
|
||||
snapshot_service = HookSnapshotService(
|
||||
plan = ExecutionPlan.build(
|
||||
self.catalog,
|
||||
selected_plugins=snapshot_selected_plugins,
|
||||
config=config,
|
||||
derived_config=derived_config,
|
||||
runtime="archivebox",
|
||||
)
|
||||
plugins = plan.plugins
|
||||
snapshot_phase_timeout = plan.snapshot_timeout + 120.0
|
||||
await plan.seed_config(self.bus, parent_event=crawl_start_event)
|
||||
snapshot_service = plan.attach_snapshot_service(
|
||||
self.bus,
|
||||
url=snapshot["url"],
|
||||
snapshot=abx_snapshot,
|
||||
output_dir=output_dir,
|
||||
plugins=plugins,
|
||||
config=RuntimeConfig(
|
||||
user=GlobalConfig(**{**config, "ABX_RUNTIME": "archivebox"}),
|
||||
derived=derived_config,
|
||||
),
|
||||
snapshot_phase_timeout=snapshot_phase_timeout,
|
||||
snapshot_cleanup_enabled=True,
|
||||
snapshot_cleanup_phase_timeout=snapshot_phase_timeout,
|
||||
snapshot_service=HookSnapshotService,
|
||||
timeout_padding=120.0,
|
||||
abort_requested=self.crawl_is_cancelled,
|
||||
selected_hooks_by_plugin=None,
|
||||
emit_discovered_snapshot_events=False,
|
||||
)
|
||||
try:
|
||||
snapshot_event = SnapshotEvent(
|
||||
@ -1176,7 +1180,7 @@ class CrawlRunner:
|
||||
# SnapshotCompletedEvent is the normal projection path, but the
|
||||
# runner is the scheduler owner. Finalize idempotently here too
|
||||
# so a completed snapshot cannot remain STARTED if the event was
|
||||
# observed before its DB projector advanced the state machine.
|
||||
# observed before its DB projector advanced the lifecycle.
|
||||
crawl_limit_stop_reason = CrawlLimitState.from_config(config).get_stop_reason()
|
||||
await sync_to_async(finalize_completed_snapshot, thread_sensitive=True)(
|
||||
snapshot_id,
|
||||
@ -1189,18 +1193,8 @@ class CrawlRunner:
|
||||
await self.enqueue_discovered_snapshots_from_outputs(snapshot)
|
||||
|
||||
def _seal_when_last_snapshot_finished() -> None:
|
||||
# run_snapshot replaces self.crawl per-snapshot, so multiple
|
||||
# concurrent tasks each load a fresh Crawl/SM pointing at the
|
||||
# same DB row. The "no open snapshots" check is non-atomic
|
||||
# with sm.seal(), so two tasks racing to finish the last
|
||||
# snapshot can both pass the guard. The first call drives
|
||||
# the SM to a final state (engine.running=False); the loser
|
||||
# then raises TransitionNotAllowed even though current_state
|
||||
# still reads STARTED off its stale model field. Re-read the
|
||||
# row right before the call and swallow the race so the
|
||||
# task that lost the lap doesn't fail the whole snapshot.
|
||||
from statemachine.exceptions import TransitionNotAllowed
|
||||
|
||||
# Re-read immediately before the idempotent conditional
|
||||
# update so concurrent last-snapshot completions are safe.
|
||||
crawl = self.crawl
|
||||
crawl.refresh_from_db(fields=["status"])
|
||||
if crawl.status != crawl.StatusChoices.STARTED:
|
||||
@ -1209,12 +1203,7 @@ class CrawlRunner:
|
||||
status__in=crawl.snapshot_set.model.OPEN_STATES,
|
||||
).exists():
|
||||
return
|
||||
try:
|
||||
crawl.sm.seal()
|
||||
except TransitionNotAllowed:
|
||||
# Another task sealed it between our refresh and the
|
||||
# SM call. Idempotent by design.
|
||||
pass
|
||||
crawl.seal()
|
||||
|
||||
await sync_to_async(_seal_when_last_snapshot_finished, thread_sensitive=True)()
|
||||
finally:
|
||||
@ -1228,7 +1217,7 @@ class CrawlRunner:
|
||||
return
|
||||
# Limit stops are runner-owned cancellation decisions, not normal
|
||||
# "all ArchiveResults finished" lifecycle seals. Updating the row
|
||||
# directly avoids racing the state machine's in-memory state while
|
||||
# directly avoids racing a concurrent lifecycle update while
|
||||
# concurrent snapshot tasks are stopping because the crawl-wide limit
|
||||
# has already been reached.
|
||||
snapshot.update_and_requeue(
|
||||
@ -1454,11 +1443,8 @@ def snapshot_hooks_for_pending_archiveresults(snapshot) -> list[tuple[str, str]]
|
||||
crawl_plugin_names = [name.strip() for name in str((snapshot.crawl.config or {}).get("PLUGINS") or "").split(",") if name.strip()]
|
||||
config_plugin_names = [name.strip() for name in str(config.PLUGINS or "").split(",") if name.strip()]
|
||||
plugin_names = snapshot_plugin_names or crawl_plugin_names or config_plugin_names or get_enabled_plugins(config=config)
|
||||
plugins = (
|
||||
filter_plugins(_discover_archivebox_plugins(), plugin_names, include_providers=True)
|
||||
if plugin_names
|
||||
else _discover_archivebox_plugins()
|
||||
)
|
||||
catalog = _discover_archivebox_catalog()
|
||||
plugins = catalog.select(plugin_names).plugins if plugin_names else catalog.plugins
|
||||
if snapshot.url == Snapshot.INTERNAL_INPUT_URL:
|
||||
plugins = {name: plugin for name, plugin in plugins.items() if getattr(plugin.config, "x_accepts_internal_input", False)}
|
||||
return sorted((plugin.name, hook.name) for plugin in plugins.values() for hook in plugin.filter_hooks("Snapshot"))
|
||||
@ -1548,7 +1534,7 @@ def _run_due_crawl_locked(crawl, *, lock_seconds: int, interactive_interrupts: b
|
||||
if not crawl.claim_processing_lock(lock_seconds=lock_seconds):
|
||||
return False
|
||||
crawl.refresh_from_db()
|
||||
crawl.sm.tick()
|
||||
crawl.advance_lifecycle()
|
||||
return True
|
||||
|
||||
# retry_at is the only queue/ownership signal the runner sees.
|
||||
@ -1583,7 +1569,7 @@ def _run_due_crawl_locked(crawl, *, lock_seconds: int, interactive_interrupts: b
|
||||
return False
|
||||
crawl.refresh_from_db()
|
||||
if crawl.status == crawl.StatusChoices.STARTED and crawl.is_finished():
|
||||
crawl.sm.tick()
|
||||
crawl.advance_lifecycle()
|
||||
return True
|
||||
_runner_console_line(crawl=crawl)
|
||||
run_crawl(str(crawl.id), process_discovered_snapshots_inline=True, interactive_interrupts=interactive_interrupts)
|
||||
@ -1593,7 +1579,7 @@ def _run_due_crawl_locked(crawl, *, lock_seconds: int, interactive_interrupts: b
|
||||
if not type(crawl).claim_for_worker(crawl, lock_seconds=lock_seconds):
|
||||
return False
|
||||
_runner_console_line(crawl=crawl, status="SEALED")
|
||||
crawl.cleanup()
|
||||
crawl.cleanup_runtime()
|
||||
crawl.update_and_requeue(retry_at=None)
|
||||
return True
|
||||
|
||||
@ -1628,7 +1614,7 @@ def _run_due_snapshot_locked(snapshot, *, lock_seconds: int, interactive_interru
|
||||
|
||||
if snapshot.is_paused:
|
||||
# Paused work never executes out of band. Preserve the lifecycle marker
|
||||
# until an explicit resume moves it through the normal state machine.
|
||||
# until an explicit resume moves it through the normal lifecycle.
|
||||
from archivebox.core.models import ArchiveResult
|
||||
|
||||
ArchiveResult.pause_queryset(snapshot.archiveresult_set.all())
|
||||
@ -1724,7 +1710,7 @@ def _run_due_snapshot_locked(snapshot, *, lock_seconds: int, interactive_interru
|
||||
# browser-uploaded rows are reused and queued so the server adds its
|
||||
# outputs to that plugin result instead of creating a sibling row.
|
||||
snapshot.create_pending_archiveresults(hooks=snapshot_hooks_for_pending_archiveresults(snapshot))
|
||||
snapshot.sm.tick()
|
||||
snapshot.advance_lifecycle()
|
||||
snapshot.refresh_from_db()
|
||||
if snapshot.status == Snapshot.StatusChoices.SEALED:
|
||||
_runner_console_line(crawl_id=snapshot.crawl_id, snapshot=snapshot, status="SEALED")
|
||||
@ -2095,7 +2081,7 @@ def run_pending_crawls(
|
||||
from archivebox.config.common import get_config
|
||||
from archivebox.crawls.models import Crawl, CrawlSchedule
|
||||
from archivebox.core.models import ArchiveResult, Snapshot
|
||||
from archivebox.plugins.hooks import discover_hooks
|
||||
from archivebox.plugins.discovery import get_enabled_plugins, get_plugin_catalog
|
||||
from archivebox.machine.models import Process
|
||||
|
||||
crawl_claim_lock_seconds = 10
|
||||
@ -2235,8 +2221,10 @@ def run_pending_crawls(
|
||||
# hydration. Refreshing here preserves mid-run config edits while using
|
||||
# the same enabled-hook discovery path that created ArchiveResult rows.
|
||||
runtime_config = get_config()
|
||||
catalog = get_plugin_catalog()
|
||||
enabled_plugins = get_enabled_plugins(config=runtime_config)
|
||||
search_plugin_names = frozenset(
|
||||
hook.parent.name for hook in discover_hooks("Snapshot", config=runtime_config) if hook.parent.name.startswith("search_backend_")
|
||||
plugin.name for plugin, _hook in catalog.hooks("Snapshot", names=enabled_plugins) if plugin.name.startswith("search_backend_")
|
||||
)
|
||||
if _run_due_queued_plugin_result(
|
||||
search_plugin_names,
|
||||
|
||||
@ -74,10 +74,10 @@ def finalize_completed_snapshot(
|
||||
)
|
||||
|
||||
if snapshot.status == Snapshot.StatusChoices.QUEUED:
|
||||
snapshot.sm.tick()
|
||||
snapshot.advance_lifecycle()
|
||||
snapshot.refresh_from_db()
|
||||
if snapshot.status == Snapshot.StatusChoices.STARTED and snapshot.is_finished_processing():
|
||||
snapshot.sm.seal()
|
||||
snapshot.seal()
|
||||
snapshot.refresh_from_db()
|
||||
|
||||
snapshot.write_index_jsonl(output_dir=output_dir)
|
||||
@ -119,7 +119,7 @@ class SnapshotService(BaseService):
|
||||
hooks = await sync_to_async(snapshot_hooks_for_pending_archiveresults, thread_sensitive=True)(snapshot)
|
||||
await sync_to_async(snapshot.create_pending_archiveresults, thread_sensitive=True)(hooks=hooks)
|
||||
try:
|
||||
await sync_to_async(snapshot.sm.tick, thread_sensitive=True)()
|
||||
await sync_to_async(snapshot.advance_lifecycle, thread_sensitive=True)()
|
||||
except ValidationError as err:
|
||||
if "ArchiveBox cannot archive its own admin, web, api, or snapshot URLs." not in str(err):
|
||||
raise
|
||||
|
||||
@ -1683,7 +1683,7 @@ def install_real_binary(
|
||||
binproviders: str = "env",
|
||||
overrides: dict[str, dict[str, Any]] | None = None,
|
||||
):
|
||||
"""Install and persist a real binary through the normal Binary state machine."""
|
||||
"""Install and persist a real binary through the normal Binary lifecycle."""
|
||||
from archivebox.machine.models import Binary, Machine
|
||||
|
||||
binary = Binary.objects.create(
|
||||
@ -1693,7 +1693,7 @@ def install_real_binary(
|
||||
overrides=overrides or {},
|
||||
status=Binary.StatusChoices.QUEUED,
|
||||
)
|
||||
assert binary.tick_claimed(lock_seconds=600)
|
||||
assert binary.install_claimed(lock_seconds=600)
|
||||
binary.refresh_from_db()
|
||||
assert binary.status == Binary.StatusChoices.INSTALLED
|
||||
assert binary.retry_at is None
|
||||
|
||||
@ -195,7 +195,7 @@ def test_crawl_pause_resume_api_cascades_archiveresults_and_leaves_finished_snap
|
||||
hook_name="on_Snapshot__93_hashes.py",
|
||||
lib_dir=lib_dir,
|
||||
)
|
||||
sealed_snapshot.sm.seal()
|
||||
sealed_snapshot.seal()
|
||||
sealed_snapshot.refresh_from_db()
|
||||
assert sealed_snapshot.status == Snapshot.StatusChoices.SEALED
|
||||
assert sealed_snapshot.retry_at is None
|
||||
|
||||
@ -554,14 +554,17 @@ def test_collect_output_metadata_preserves_file_metadata():
|
||||
|
||||
|
||||
def test_collect_output_metadata_detects_warc_gz_mimetype(tmp_path):
|
||||
from archivebox.services.archive_result_service import _collect_output_metadata
|
||||
from abx_dl.output_files import OutputManifest
|
||||
|
||||
plugin_dir = tmp_path / "wget"
|
||||
warc_file = plugin_dir / "warc" / "capture.warc.gz"
|
||||
warc_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
warc_file.write_bytes(b"warc-bytes")
|
||||
|
||||
output_files, output_size, output_mimetypes = _collect_output_metadata(plugin_dir)
|
||||
manifest = OutputManifest.scan(plugin_dir)
|
||||
output_files = manifest.as_mapping()
|
||||
output_size = manifest.total_size
|
||||
output_mimetypes = ",".join(manifest.mimetypes)
|
||||
|
||||
assert output_files["warc/capture.warc.gz"] == {
|
||||
"extension": "gz",
|
||||
|
||||
@ -22,13 +22,13 @@ def _runtime_env(data_dir: Path, *, lib_dir: Path | None = None, **extra: str) -
|
||||
}
|
||||
|
||||
|
||||
def _run_real_binary_state_machine(data_dir: Path, *, name: str, binproviders: str, env: dict[str, str]):
|
||||
"""Run a real Binary model through its abxpkg-backed state machine."""
|
||||
def _run_real_binary_lifecycle(data_dir: Path, *, name: str, binproviders: str, env: dict[str, str]):
|
||||
"""Run a real Binary model through its abxpkg-backed lifecycle."""
|
||||
script = (
|
||||
"from archivebox.machine.models import Binary, Machine; "
|
||||
f"binary = Binary.objects.create(machine=Machine.current(), name={name!r}, binproviders={binproviders!r}, status=Binary.StatusChoices.QUEUED); "
|
||||
"assert binary.tick_claimed(lock_seconds=600); "
|
||||
"print('BINARY_STATE_MACHINE_E2E_DONE')"
|
||||
"assert binary.install_claimed(lock_seconds=600); "
|
||||
"print('BINARY_LIFECYCLE_E2E_DONE')"
|
||||
)
|
||||
return run_archivebox_cmd(
|
||||
["shell", "-c", script],
|
||||
@ -59,7 +59,7 @@ def test_binary_request_preserves_native_overrides_in_db():
|
||||
status=Binary.StatusChoices.QUEUED,
|
||||
retry_at=timezone.now(),
|
||||
)
|
||||
assert binary.tick_claimed(lock_seconds=600)
|
||||
assert binary.install_claimed(lock_seconds=600)
|
||||
binary.refresh_from_db()
|
||||
assert binary.status == Binary.StatusChoices.INSTALLED
|
||||
assert Path(binary.abspath).resolve() == Path(sys.executable).resolve()
|
||||
@ -98,11 +98,11 @@ def test_binary_request_installs_env_binary_and_recovers_stale_cache(initialized
|
||||
host_binary = shutil.which(name)
|
||||
assert host_binary is not None
|
||||
runtime_env = _runtime_env(initialized_archive)
|
||||
_cmd_result = _run_real_binary_state_machine(initialized_archive, name=name, binproviders="env", env=runtime_env)
|
||||
_cmd_result = _run_real_binary_lifecycle(initialized_archive, name=name, binproviders="env", env=runtime_env)
|
||||
stdout, stderr, returncode = _cmd_result.stdout, _cmd_result.stderr, _cmd_result.returncode
|
||||
|
||||
assert returncode == 0, stderr
|
||||
assert "BINARY_STATE_MACHINE_E2E_DONE" in stdout
|
||||
assert "BINARY_LIFECYCLE_E2E_DONE" in stdout
|
||||
|
||||
with use_archivebox_db(initialized_archive):
|
||||
binary = Binary.objects.get(name=name)
|
||||
@ -203,7 +203,7 @@ def test_missing_binary_request_stays_queued_then_recovers_when_provider_can_res
|
||||
provider_bin_dir = initialized_archive / "lib" / "pip" / "venv" / "bin"
|
||||
runtime_env = _runtime_env(initialized_archive)
|
||||
|
||||
_cmd_result = _run_real_binary_state_machine(initialized_archive, name=name, binproviders="env", env=runtime_env)
|
||||
_cmd_result = _run_real_binary_lifecycle(initialized_archive, name=name, binproviders="env", env=runtime_env)
|
||||
stdout, stderr, returncode = _cmd_result.stdout, _cmd_result.stderr, _cmd_result.returncode
|
||||
|
||||
assert returncode != 0, stdout + stderr
|
||||
|
||||
@ -2121,7 +2121,7 @@ class TestRunDueCrawlState:
|
||||
retry_at=timezone.now(),
|
||||
)
|
||||
|
||||
snapshot.sm.tick()
|
||||
snapshot.advance_lifecycle()
|
||||
snapshot.refresh_from_db()
|
||||
|
||||
assert snapshot.status == Snapshot.StatusChoices.STARTED
|
||||
@ -3007,7 +3007,7 @@ class TestRecoverOrchestratorStateRedFailureModes:
|
||||
modified_at=now + timedelta(seconds=1),
|
||||
)
|
||||
|
||||
snapshot.sm.seal()
|
||||
snapshot.seal()
|
||||
snapshot.refresh_from_db()
|
||||
assert snapshot.status == Snapshot.StatusChoices.SEALED
|
||||
assert snapshot.retry_at is None
|
||||
|
||||
@ -261,7 +261,9 @@ def test_snapshot_started_state_keeps_retry_at_lease():
|
||||
retry_at=before,
|
||||
)
|
||||
|
||||
assert snapshot.tick_claimed(lock_seconds=60) is True
|
||||
assert snapshot.claim_processing_lock(lock_seconds=60) is True
|
||||
snapshot.refresh_from_db()
|
||||
assert snapshot.advance_lifecycle() is True
|
||||
|
||||
snapshot.refresh_from_db()
|
||||
assert snapshot.status == Snapshot.StatusChoices.STARTED
|
||||
|
||||
@ -4,8 +4,8 @@ Unit tests for machine module models: Machine, NetworkInterface, Binary, Process
|
||||
Tests cover:
|
||||
1. Machine model creation and current() method
|
||||
2. NetworkInterface model and network detection
|
||||
3. Binary model lifecycle and state machine
|
||||
4. Process model lifecycle, hierarchy, and state machine
|
||||
3. Binary model lifecycle
|
||||
4. Process model lifecycle and hierarchy
|
||||
5. JSONL serialization/deserialization
|
||||
6. Manager methods
|
||||
7. Process tracking methods (replacing pid_utils)
|
||||
@ -29,8 +29,6 @@ from archivebox.machine.models import (
|
||||
NetworkInterface,
|
||||
Binary,
|
||||
Process,
|
||||
BinaryMachine,
|
||||
ProcessMachine,
|
||||
MACHINE_RECHECK_INTERVAL,
|
||||
PID_REUSE_WINDOW,
|
||||
PROCESS_TIMEOUT_GRACE,
|
||||
@ -509,27 +507,19 @@ class TestBinaryModel:
|
||||
assert symlink.resolve() == source.resolve()
|
||||
|
||||
|
||||
class TestBinaryStateMachine:
|
||||
"""Test the BinaryMachine state machine."""
|
||||
class TestBinaryLifecycle:
|
||||
"""Test Binary lifecycle prerequisites."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup_binary(self, binary):
|
||||
self.binary = binary
|
||||
|
||||
def test_binary_state_machine_initial_state(self):
|
||||
"""BinaryMachine should start in queued state."""
|
||||
sm = BinaryMachine(self.binary)
|
||||
assert sm.current_state_value == Binary.StatusChoices.QUEUED
|
||||
|
||||
def test_binary_state_machine_can_start(self):
|
||||
"""BinaryMachine.can_start() should check name and binproviders."""
|
||||
sm = BinaryMachine(self.binary)
|
||||
assert sm.can_install()
|
||||
def test_binary_can_install_checks_name_and_binproviders(self):
|
||||
assert self.binary.can_install
|
||||
|
||||
self.binary.binproviders = ""
|
||||
self.binary.save()
|
||||
sm = BinaryMachine(self.binary)
|
||||
assert not sm.can_install()
|
||||
assert not self.binary.can_install
|
||||
|
||||
|
||||
class TestProcessModel:
|
||||
@ -1002,38 +992,5 @@ class TestProcessClassMethods:
|
||||
assert child.exit_code == 143
|
||||
|
||||
|
||||
class TestProcessStateMachine:
|
||||
"""Test the ProcessMachine state machine."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup_process(self, process):
|
||||
self.process = process
|
||||
|
||||
def test_process_state_machine_initial_state(self):
|
||||
"""ProcessMachine should start in queued state."""
|
||||
sm = ProcessMachine(self.process)
|
||||
assert sm.current_state_value == Process.StatusChoices.QUEUED
|
||||
|
||||
def test_process_state_machine_can_start(self):
|
||||
"""ProcessMachine.can_start() should check cmd and machine."""
|
||||
sm = ProcessMachine(self.process)
|
||||
assert sm.can_start()
|
||||
|
||||
self.process.cmd = []
|
||||
self.process.save()
|
||||
sm = ProcessMachine(self.process)
|
||||
assert not sm.can_start()
|
||||
|
||||
def test_process_state_machine_is_exited(self):
|
||||
"""ProcessMachine.is_exited() should check exit_code."""
|
||||
sm = ProcessMachine(self.process)
|
||||
assert not sm.is_exited()
|
||||
|
||||
self.process.exit_code = 0
|
||||
self.process.save()
|
||||
sm = ProcessMachine(self.process)
|
||||
assert sm.is_exited()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
@ -6,7 +6,7 @@ Migration tests from 0.8.x to 0.9.x.
|
||||
- Crawl model for grouping URLs
|
||||
- Seed model (removed in 0.9.x)
|
||||
- UUID primary keys for Snapshot
|
||||
- Status fields for state machine
|
||||
- Status fields for queued lifecycle processing
|
||||
- New fields like depth, retry_at, etc.
|
||||
"""
|
||||
|
||||
|
||||
@ -3,22 +3,15 @@ __package__ = "archivebox.workers"
|
||||
import inspect
|
||||
import logging
|
||||
|
||||
from typing import Any, ClassVar, Protocol, cast
|
||||
from collections.abc import Iterable
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
from statemachine.callbacks import SPECS_ALL
|
||||
from statemachine.dispatcher import Listener, Listeners
|
||||
from statemachine.graph import iterate_states_and_transitions
|
||||
from typing import Any, ClassVar, cast
|
||||
|
||||
from django.db import models
|
||||
from django.core import checks
|
||||
from django.utils import timezone
|
||||
from django.utils.functional import classproperty
|
||||
from django_stubs_ext.db.models import TypedModelMeta
|
||||
|
||||
from statemachine import registry, StateMachine, State
|
||||
|
||||
|
||||
class DefaultStatusChoices(models.TextChoices):
|
||||
QUEUED = "queued", "Queued"
|
||||
@ -43,230 +36,59 @@ MODULE_PATH = Path(__file__).resolve()
|
||||
REPO_ROOT = MODULE_PATH.parents[2]
|
||||
PACKAGE_ROOT = MODULE_PATH.parents[1]
|
||||
|
||||
ObjectState = State | str
|
||||
ObjectStateList = Iterable[ObjectState]
|
||||
|
||||
class ModelWithQueue(models.Model):
|
||||
"""Durable queue fields and atomic lease operations shared by work rows.
|
||||
|
||||
class ModelStateMachine(Protocol):
|
||||
def tick(self) -> Any: ...
|
||||
Concrete models own lifecycle transitions. This mixin only owns the common
|
||||
database queue protocol: status, retry_at, pause/resume, and claims.
|
||||
"""
|
||||
|
||||
def pause_requested(self) -> Any: ...
|
||||
|
||||
def resume_requested(self) -> Any: ...
|
||||
|
||||
|
||||
class BaseModelWithStateMachine(models.Model):
|
||||
StatusChoices: ClassVar[type[DefaultStatusChoices]]
|
||||
|
||||
# status: models.CharField
|
||||
# retry_at: models.DateTimeField
|
||||
|
||||
state_machine_name: str | None = None
|
||||
state_field_name: str
|
||||
state_machine_attr: str = "sm"
|
||||
bind_events_as_methods: bool = False
|
||||
StatusChoices: ClassVar[type[models.TextChoices]] = DefaultStatusChoices
|
||||
INITIAL_STATE: ClassVar[str] = DefaultStatusChoices.QUEUED
|
||||
ACTIVE_STATE: ClassVar[str] = DefaultStatusChoices.STARTED
|
||||
FINAL_STATES: ClassVar[tuple[str, ...]] = (DefaultStatusChoices.SEALED,)
|
||||
warn_on_save_outside_runner: ClassVar[bool] = True
|
||||
|
||||
active_state: ObjectState
|
||||
retry_at_field_name: str
|
||||
status: models.CharField = models.CharField(**default_status_field.deconstruct()[3])
|
||||
retry_at: models.DateTimeField = models.DateTimeField(**default_retry_at_field.deconstruct()[3])
|
||||
|
||||
class Meta(TypedModelMeta):
|
||||
app_label = "workers"
|
||||
abstract = True
|
||||
|
||||
@property
|
||||
def sm(self) -> StateMachine:
|
||||
"""Build the python-statemachine wrapper only at transition callsites.
|
||||
|
||||
This model is loaded by high-volume paths that do not drive lifecycle
|
||||
transitions: admin lists, progress polling, index-only maintenance, and
|
||||
bulk recovery scans all instantiate thousands of rows just to read or
|
||||
update ordinary columns. python-statemachine setup is correct but not
|
||||
free: it creates per-instance state wrappers, callback registries,
|
||||
queues, locks, and callback adapters. Paying that cost from Django's
|
||||
model __init__ made plain ORM materialization scale with state-machine
|
||||
setup instead of row decoding.
|
||||
|
||||
ArchiveBox drives lifecycle transitions explicitly through `.sm`
|
||||
(`snapshot.sm.tick()`, `crawl.sm.seal()`, etc.), so the machine can be
|
||||
cached on first use without changing the state model. Code that only
|
||||
needs database fields never constructs one.
|
||||
"""
|
||||
try:
|
||||
machine = vars(self)["_archivebox_state_machine"]
|
||||
except KeyError:
|
||||
machine = self.StateMachineClass(self, state_field=self.state_field_name)
|
||||
vars(self)["_archivebox_state_machine"] = machine
|
||||
return cast(StateMachine, machine)
|
||||
FINAL_OR_ACTIVE_STATES: ClassVar[tuple[str, ...]] = (*FINAL_STATES, ACTIVE_STATE)
|
||||
|
||||
@classmethod
|
||||
def status_counts(cls, queryset: models.QuerySet | None = None, statuses: Iterable[str] | None = None) -> dict[str, int]:
|
||||
"""Count requested statuses with separate indexed COUNT probes.
|
||||
|
||||
For live/progress views this is often faster on large SQLite data dirs
|
||||
than a grouped aggregate, because each status can use the status index
|
||||
directly and the caller usually needs only a few states.
|
||||
"""
|
||||
qs = queryset if queryset is not None else cls.objects.all()
|
||||
return {status: qs.filter(status=status).count() for status in (statuses or cls.StatusChoices.values)}
|
||||
|
||||
@classmethod
|
||||
def check(cls, sender=None, **kwargs):
|
||||
import sys
|
||||
|
||||
# Skip state machine checks during makemigrations to avoid premature registry access
|
||||
if "makemigrations" in sys.argv:
|
||||
return super().check(**kwargs)
|
||||
|
||||
errors = super().check(**kwargs)
|
||||
|
||||
found_id_field = False
|
||||
found_status_field = False
|
||||
found_retry_at_field = False
|
||||
|
||||
for field in cls._meta.get_fields():
|
||||
if getattr(field, "_is_state_field", False):
|
||||
if cls.state_field_name == field.name:
|
||||
found_status_field = True
|
||||
if getattr(field, "choices", None) != cls.StatusChoices.choices:
|
||||
errors.append(
|
||||
checks.Error(
|
||||
f"{cls.__name__}.{field.name} must have choices set to {cls.__name__}.StatusChoices.choices",
|
||||
hint=f"{cls.__name__}.{field.name}.choices = {getattr(field, 'choices', None)!r}",
|
||||
obj=cls,
|
||||
id="workers.E011",
|
||||
),
|
||||
)
|
||||
if getattr(field, "_is_retry_at_field", False):
|
||||
if cls.retry_at_field_name == field.name:
|
||||
found_retry_at_field = True
|
||||
if field.name == "id" and getattr(field, "primary_key", False):
|
||||
found_id_field = True
|
||||
|
||||
if not found_status_field:
|
||||
errors.append(
|
||||
checks.Error(
|
||||
f"{cls.__name__}.state_field_name must be defined and point to a StatusField()",
|
||||
hint=f"{cls.__name__}.state_field_name = {cls.state_field_name!r} but {cls.__name__}.{cls.state_field_name!r} was not found or does not refer to StatusField",
|
||||
obj=cls,
|
||||
id="workers.E012",
|
||||
),
|
||||
)
|
||||
if not found_retry_at_field:
|
||||
errors.append(
|
||||
checks.Error(
|
||||
f"{cls.__name__}.retry_at_field_name must be defined and point to a RetryAtField()",
|
||||
hint=f"{cls.__name__}.retry_at_field_name = {cls.retry_at_field_name!r} but {cls.__name__}.{cls.retry_at_field_name!r} was not found or does not refer to RetryAtField",
|
||||
obj=cls,
|
||||
id="workers.E013",
|
||||
),
|
||||
)
|
||||
|
||||
if not found_id_field:
|
||||
errors.append(
|
||||
checks.Error(
|
||||
f"{cls.__name__} must have an id field that is a primary key",
|
||||
hint=f"{cls.__name__}.id field missing or not configured as primary key",
|
||||
obj=cls,
|
||||
id="workers.E014",
|
||||
),
|
||||
)
|
||||
|
||||
if not isinstance(cls.state_machine_name, str):
|
||||
errors.append(
|
||||
checks.Error(
|
||||
f"{cls.__name__}.state_machine_name must be a dotted-import path to a StateMachine class",
|
||||
hint=f"{cls.__name__}.state_machine_name = {cls.state_machine_name!r}",
|
||||
obj=cls,
|
||||
id="workers.E015",
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
cls.StateMachineClass
|
||||
except Exception as err:
|
||||
errors.append(
|
||||
checks.Error(
|
||||
f"{cls.__name__}.state_machine_name must point to a valid StateMachine class, but got {type(err).__name__} {err} when trying to access {cls.__name__}.StateMachineClass",
|
||||
hint=f"{cls.__name__}.state_machine_name = {cls.state_machine_name!r}",
|
||||
obj=cls,
|
||||
id="workers.E016",
|
||||
),
|
||||
)
|
||||
|
||||
if cls.INITIAL_STATE not in cls.StatusChoices.values:
|
||||
errors.append(
|
||||
checks.Error(
|
||||
f"{cls.__name__}.StateMachineClass.initial_state must be present within {cls.__name__}.StatusChoices",
|
||||
hint=f"{cls.__name__}.StateMachineClass.initial_state = {cls.StateMachineClass.initial_state!r}",
|
||||
obj=cls,
|
||||
id="workers.E017",
|
||||
),
|
||||
)
|
||||
|
||||
if cls.ACTIVE_STATE not in cls.StatusChoices.values:
|
||||
errors.append(
|
||||
checks.Error(
|
||||
f"{cls.__name__}.active_state must be set to a valid State present within {cls.__name__}.StatusChoices",
|
||||
hint=f"{cls.__name__}.active_state = {cls.active_state!r}",
|
||||
obj=cls,
|
||||
id="workers.E018",
|
||||
),
|
||||
)
|
||||
|
||||
for state in cls.FINAL_STATES:
|
||||
if state not in cls.StatusChoices.values:
|
||||
errors.append(
|
||||
checks.Error(
|
||||
f"{cls.__name__}.StateMachineClass.final_states must all be present within {cls.__name__}.StatusChoices",
|
||||
hint=f"{cls.__name__}.StateMachineClass.final_states = {cls.StateMachineClass.final_states!r}",
|
||||
obj=cls,
|
||||
id="workers.E019",
|
||||
),
|
||||
)
|
||||
break
|
||||
return errors
|
||||
|
||||
@staticmethod
|
||||
def _state_to_str(state: ObjectState) -> str:
|
||||
"""Convert a statemachine.State, models.TextChoices.choices value, or Enum value to a str"""
|
||||
return str(state.value) if isinstance(state, State) else str(state)
|
||||
|
||||
@property
|
||||
def RETRY_AT(self) -> datetime:
|
||||
return getattr(self, self.retry_at_field_name)
|
||||
def RETRY_AT(self) -> datetime | None:
|
||||
return self.retry_at
|
||||
|
||||
@RETRY_AT.setter
|
||||
def RETRY_AT(self, value: datetime):
|
||||
setattr(self, self.retry_at_field_name, value)
|
||||
def RETRY_AT(self, value: datetime | None) -> None:
|
||||
self.retry_at = value
|
||||
|
||||
@property
|
||||
def STATE(self) -> str:
|
||||
return getattr(self, self.state_field_name)
|
||||
return self.status
|
||||
|
||||
@STATE.setter
|
||||
def STATE(self, value: str):
|
||||
setattr(self, self.state_field_name, value)
|
||||
def STATE(self, value: str) -> None:
|
||||
self.status = value
|
||||
|
||||
def bump_retry_at(self, seconds: int = 10):
|
||||
self.RETRY_AT = timezone.now() + timedelta(seconds=seconds)
|
||||
def bump_retry_at(self, seconds: int = 10) -> None:
|
||||
self.retry_at = timezone.now() + timedelta(seconds=seconds)
|
||||
|
||||
@property
|
||||
def is_paused(self) -> bool:
|
||||
paused_state = getattr(self.StatusChoices, "PAUSED", None)
|
||||
return paused_state is not None and self.STATE == paused_state
|
||||
return paused_state is not None and self.status == paused_state
|
||||
|
||||
def safe_update(self, update_fields: dict[str, Any], *, refresh: bool = True, extra_filter: dict[str, Any] | None = None) -> bool:
|
||||
"""
|
||||
Atomic single-row UPDATE for scheduler writes that bypass save().
|
||||
|
||||
The write is unconditional unless the caller passes extra_filter — the
|
||||
previous implicit modified_at CAS predicate spuriously collided with
|
||||
concurrent writers to unrelated fields (every save bumps modified_at),
|
||||
which silently dropped state-machine transitions. Callers that need a
|
||||
transition guard (only advance from state A to state B; only requeue a
|
||||
row still holding lease X) pass extra_filter explicitly.
|
||||
"""
|
||||
values = dict(update_fields)
|
||||
values.setdefault("modified_at", timezone.now())
|
||||
queryset = type(self).objects.filter(pk=self.pk)
|
||||
@ -274,17 +96,14 @@ class BaseModelWithStateMachine(models.Model):
|
||||
queryset = queryset.filter(**extra_filter)
|
||||
updated = queryset.update(**values)
|
||||
if updated != 1 and extra_filter:
|
||||
current = type(self).objects.filter(pk=self.pk).values(self.state_field_name).first()
|
||||
current_status = current.get(self.state_field_name) if current else "<deleted>"
|
||||
current = type(self).objects.filter(pk=self.pk).values("status").first()
|
||||
logger.info(
|
||||
"SafeUpdateGuardMiss: %s row %s extra_filter=%s did not match (current %s=%s, loaded %s=%s); update_fields=%s skipped",
|
||||
"SafeUpdateGuardMiss: %s row %s extra_filter=%s current_status=%s loaded_status=%s update_fields=%s skipped",
|
||||
type(self).__name__,
|
||||
self.pk,
|
||||
extra_filter,
|
||||
self.state_field_name,
|
||||
current_status,
|
||||
self.state_field_name,
|
||||
self.STATE,
|
||||
current.get("status") if current else "<deleted>",
|
||||
self.status,
|
||||
sorted(values),
|
||||
)
|
||||
if refresh:
|
||||
@ -294,7 +113,7 @@ class BaseModelWithStateMachine(models.Model):
|
||||
pass
|
||||
return updated == 1
|
||||
|
||||
def save(self, *args, **kwargs):
|
||||
def save(self, *args: Any, **kwargs: Any) -> None:
|
||||
from archivebox.machine.models import Process
|
||||
|
||||
process = Process.current()
|
||||
@ -325,12 +144,11 @@ class BaseModelWithStateMachine(models.Model):
|
||||
finally:
|
||||
del frame
|
||||
logger.warning(
|
||||
"%s.save() outside runner process: id=%s status=%s retry_at=%s process=%s root=%s caller=%s; "
|
||||
"queue/status writes outside the runner should usually use safe_update()",
|
||||
"%s.save() outside runner process: id=%s status=%s retry_at=%s process=%s root=%s caller=%s",
|
||||
type(self).__name__,
|
||||
self.pk,
|
||||
self.STATE,
|
||||
self.RETRY_AT,
|
||||
self.status,
|
||||
self.retry_at,
|
||||
process.process_type,
|
||||
root_type,
|
||||
caller,
|
||||
@ -338,327 +156,74 @@ class BaseModelWithStateMachine(models.Model):
|
||||
super().save(*args, **kwargs)
|
||||
|
||||
def pause(self, *, save: bool = True) -> bool:
|
||||
try:
|
||||
paused_state = self.StatusChoices.PAUSED
|
||||
except AttributeError:
|
||||
return False
|
||||
if self.STATE in self.FINAL_STATES or self.is_paused:
|
||||
paused_state = getattr(self.StatusChoices, "PAUSED", None)
|
||||
if paused_state is None or self.status in self.FINAL_STATES or self.is_paused:
|
||||
return False
|
||||
previous_status = self.status
|
||||
self.status = paused_state
|
||||
self.retry_at = RETRY_AT_MAX
|
||||
if save:
|
||||
cast(ModelStateMachine, self.sm).pause_requested()
|
||||
self.refresh_from_db()
|
||||
return self.is_paused
|
||||
self.STATE = paused_state
|
||||
self.RETRY_AT = RETRY_AT_MAX
|
||||
return self.safe_update(
|
||||
{"status": paused_state, "retry_at": RETRY_AT_MAX},
|
||||
extra_filter={"status": previous_status},
|
||||
)
|
||||
return True
|
||||
|
||||
def resume(self, *, when: datetime | None = None, save: bool = True) -> bool:
|
||||
try:
|
||||
paused_state = self.StatusChoices.PAUSED
|
||||
except AttributeError:
|
||||
return False
|
||||
if not self.is_paused:
|
||||
paused_state = getattr(self.StatusChoices, "PAUSED", None)
|
||||
if paused_state is None or not self.is_paused:
|
||||
return False
|
||||
resume_at = when or timezone.now()
|
||||
self.status = self.StatusChoices.QUEUED
|
||||
self.retry_at = resume_at
|
||||
if save:
|
||||
if when is None:
|
||||
cast(ModelStateMachine, self.sm).resume_requested()
|
||||
self.refresh_from_db()
|
||||
return self.STATE == self.StatusChoices.QUEUED
|
||||
self.STATE = self.StatusChoices.QUEUED
|
||||
self.RETRY_AT = when or timezone.now()
|
||||
updated = self.safe_update(
|
||||
{
|
||||
self.state_field_name: self.StatusChoices.QUEUED,
|
||||
self.retry_at_field_name: self.RETRY_AT,
|
||||
},
|
||||
extra_filter={self.state_field_name: paused_state},
|
||||
return self.safe_update(
|
||||
{"status": self.StatusChoices.QUEUED, "retry_at": resume_at},
|
||||
extra_filter={"status": paused_state},
|
||||
)
|
||||
return updated
|
||||
self.STATE = self.StatusChoices.QUEUED
|
||||
self.RETRY_AT = when or timezone.now()
|
||||
return True
|
||||
|
||||
def update_and_requeue(self, *, refresh: bool = True, **kwargs) -> bool:
|
||||
"""
|
||||
Scheduler-facing wrapper around safe_update().
|
||||
|
||||
Call this when a state-machine row should become visible to the
|
||||
runner. It preserves the current retry_at lease as an additional guard
|
||||
while safe_update() owns the modified_at CAS write and refresh.
|
||||
"""
|
||||
# retry_at is the scheduler lease, but it is not enough by itself:
|
||||
# sealed maintenance rows can legitimately keep the same retry_at while
|
||||
# other fields change. Include modified_at as a cheap compare-and-swap
|
||||
# guard so iterator/recovery scans never overwrite a row that the
|
||||
# runner touched after the object was read.
|
||||
current_retry_at = self.RETRY_AT
|
||||
return self.safe_update(
|
||||
dict(kwargs),
|
||||
refresh=refresh,
|
||||
extra_filter={self.retry_at_field_name: current_retry_at},
|
||||
)
|
||||
def update_and_requeue(self, *, refresh: bool = True, **kwargs: Any) -> bool:
|
||||
return self.safe_update(dict(kwargs), refresh=refresh, extra_filter={"retry_at": self.retry_at})
|
||||
|
||||
@classmethod
|
||||
def get_queue(cls):
|
||||
"""
|
||||
Get the sorted and filtered QuerySet of objects that are ready for processing.
|
||||
retry_at is the only scheduler signal; callers branch on status after selection.
|
||||
"""
|
||||
return cls.objects.filter(
|
||||
retry_at__lte=timezone.now(),
|
||||
).order_by("retry_at")
|
||||
return cls.objects.filter(retry_at__lte=timezone.now()).order_by("retry_at")
|
||||
|
||||
@classmethod
|
||||
def claim_for_worker(cls, obj: "BaseModelWithStateMachine", lock_seconds: int = 60) -> bool:
|
||||
"""
|
||||
Atomically claim a due object for processing using retry_at as the lock.
|
||||
|
||||
Correct lifecycle for any state-machine-driven work item:
|
||||
1. Queue the item by setting retry_at <= now
|
||||
2. Exactly one owner claims it by moving retry_at into the future
|
||||
3. Only that owner may call .sm.tick() and perform side effects
|
||||
4. State-machine callbacks update retry_at again when the work completes,
|
||||
backs off, or is re-queued
|
||||
|
||||
The critical rule is that future retry_at values are already owned.
|
||||
Callers must never "steal" those future timestamps and start another
|
||||
copy of the same work. That is what prevents duplicate installs, hook
|
||||
runs, and other concurrent side effects.
|
||||
|
||||
Returns True if successfully claimed, False if another worker got it
|
||||
first or the object is not currently due.
|
||||
"""
|
||||
def claim_for_worker(cls, obj: "ModelWithQueue", lock_seconds: int = 60) -> bool:
|
||||
now = timezone.now()
|
||||
lock_until = now + timedelta(seconds=lock_seconds)
|
||||
updated = cls.objects.filter(
|
||||
pk=obj.pk,
|
||||
retry_at=obj.RETRY_AT,
|
||||
retry_at__lte=now,
|
||||
).update(
|
||||
updated = cls.objects.filter(pk=obj.pk, retry_at=obj.retry_at, retry_at__lte=now).update(
|
||||
retry_at=lock_until,
|
||||
modified_at=now,
|
||||
)
|
||||
if updated == 1:
|
||||
obj.RETRY_AT = lock_until
|
||||
obj.retry_at = lock_until
|
||||
cast(Any, obj).modified_at = now
|
||||
return updated == 1
|
||||
|
||||
def claim_processing_lock(self, lock_seconds: int = 60) -> bool:
|
||||
"""
|
||||
Claim this model instance immediately before executing one state-machine tick.
|
||||
|
||||
This helper is the safe entrypoint for any direct state-machine driver
|
||||
(workers, synchronous crawl dependency installers, one-off CLI helpers).
|
||||
Calling `.sm.tick()` without claiming first turns retry_at into "just a
|
||||
schedule" instead of the ownership lock it is meant to be.
|
||||
|
||||
Returns True only for the caller that successfully moved retry_at into
|
||||
the future. False means another process already owns the work item or it
|
||||
is not currently due.
|
||||
"""
|
||||
if self.STATE in self.FINAL_STATES:
|
||||
if self.status in self.FINAL_STATES or self.retry_at is None:
|
||||
return False
|
||||
if self.RETRY_AT is None:
|
||||
return False
|
||||
|
||||
claimed = type(self).claim_for_worker(self, lock_seconds=lock_seconds)
|
||||
return claimed
|
||||
|
||||
def tick_claimed(self, lock_seconds: int = 60) -> bool:
|
||||
"""
|
||||
Claim ownership via retry_at and then execute exactly one `.sm.tick()`.
|
||||
|
||||
Future maintainers should prefer this helper over calling `.sm.tick()`
|
||||
directly whenever there is any chance another process could see the same
|
||||
queued row. If this method returns False, someone else already owns the
|
||||
work and the caller must not run side effects for it.
|
||||
"""
|
||||
if not self.claim_processing_lock(lock_seconds=lock_seconds):
|
||||
return False
|
||||
|
||||
cast(ModelStateMachine, self.sm).tick()
|
||||
self.refresh_from_db()
|
||||
return True
|
||||
|
||||
@classproperty
|
||||
def ACTIVE_STATE(cls) -> str:
|
||||
return cls._state_to_str(cls.active_state)
|
||||
|
||||
@classproperty
|
||||
def INITIAL_STATE(cls) -> str:
|
||||
initial_state = cls.StateMachineClass.initial_state
|
||||
if initial_state is None:
|
||||
raise ValueError("StateMachineClass.initial_state must not be None")
|
||||
return cls._state_to_str(initial_state)
|
||||
|
||||
@classproperty
|
||||
def FINAL_STATES(cls) -> list[str]:
|
||||
return [cls._state_to_str(state) for state in cls.StateMachineClass.final_states]
|
||||
|
||||
@classproperty
|
||||
def FINAL_OR_ACTIVE_STATES(cls) -> list[str]:
|
||||
return [*cls.FINAL_STATES, cls.ACTIVE_STATE]
|
||||
return type(self).claim_for_worker(self, lock_seconds=lock_seconds)
|
||||
|
||||
@classmethod
|
||||
def extend_choices(cls, base_choices: type[models.TextChoices]):
|
||||
"""
|
||||
Decorator to extend the base choices with extra choices, e.g.:
|
||||
|
||||
class MyModel(ModelWithStateMachine):
|
||||
|
||||
@ModelWithStateMachine.extend_choices(ModelWithStateMachine.StatusChoices)
|
||||
class StatusChoices(models.TextChoices):
|
||||
SUCCEEDED = 'succeeded'
|
||||
FAILED = 'failed'
|
||||
SKIPPED = 'skipped'
|
||||
"""
|
||||
assert issubclass(base_choices, models.TextChoices), (
|
||||
f"@extend_choices(base_choices) must be a TextChoices class, not {base_choices.__name__}"
|
||||
)
|
||||
assert issubclass(base_choices, models.TextChoices)
|
||||
|
||||
def wrapper(extra_choices: type[models.TextChoices]) -> type[models.TextChoices]:
|
||||
joined = {}
|
||||
for item in base_choices.choices:
|
||||
joined[item[0]] = item[1]
|
||||
for item in extra_choices.choices:
|
||||
joined[item[0]] = item[1]
|
||||
joined_choices = models.TextChoices("StatusChoices", joined)
|
||||
assert isinstance(joined_choices, type)
|
||||
return joined_choices
|
||||
joined = {value: label for value, label in (*base_choices.choices, *extra_choices.choices)}
|
||||
choices = models.TextChoices("StatusChoices", joined)
|
||||
assert isinstance(choices, type)
|
||||
return choices
|
||||
|
||||
return wrapper
|
||||
|
||||
@classmethod
|
||||
def StatusField(cls, **kwargs) -> models.CharField:
|
||||
"""
|
||||
Used on subclasses to extend/modify the status field with updated kwargs. e.g.:
|
||||
|
||||
class MyModel(ModelWithStateMachine):
|
||||
class StatusChoices(ModelWithStateMachine.StatusChoices):
|
||||
QUEUED = 'queued', 'Queued'
|
||||
STARTED = 'started', 'Started'
|
||||
SEALED = 'sealed', 'Sealed'
|
||||
BACKOFF = 'backoff', 'Backoff'
|
||||
FAILED = 'failed', 'Failed'
|
||||
SKIPPED = 'skipped', 'Skipped'
|
||||
|
||||
status = ModelWithStateMachine.StatusField(choices=StatusChoices.choices, default=StatusChoices.QUEUED)
|
||||
"""
|
||||
default_kwargs = default_status_field.deconstruct()[3]
|
||||
updated_kwargs = {**default_kwargs, **kwargs}
|
||||
field = models.CharField(**updated_kwargs)
|
||||
field._is_state_field = True # type: ignore
|
||||
return field
|
||||
def StatusField(cls, **kwargs: Any) -> models.CharField:
|
||||
return models.CharField(**{**default_status_field.deconstruct()[3], **kwargs})
|
||||
|
||||
@classmethod
|
||||
def RetryAtField(cls, **kwargs) -> models.DateTimeField:
|
||||
"""
|
||||
Used on subclasses to extend/modify the retry_at field with updated kwargs. e.g.:
|
||||
|
||||
class MyModel(ModelWithStateMachine):
|
||||
retry_at = ModelWithStateMachine.RetryAtField(editable=False)
|
||||
"""
|
||||
default_kwargs = default_retry_at_field.deconstruct()[3]
|
||||
updated_kwargs = {**default_kwargs, **kwargs}
|
||||
field = models.DateTimeField(**updated_kwargs)
|
||||
field._is_retry_at_field = True # type: ignore
|
||||
return field
|
||||
|
||||
@classproperty
|
||||
def StateMachineClass(cls) -> type[StateMachine]:
|
||||
"""Get the StateMachine class for the given django Model."""
|
||||
|
||||
model_state_machine_name = cls.state_machine_name
|
||||
if model_state_machine_name:
|
||||
StateMachineCls = registry.get_machine_cls(model_state_machine_name)
|
||||
assert issubclass(StateMachineCls, StateMachine)
|
||||
return StateMachineCls
|
||||
raise NotImplementedError("ActorType must define .state_machine_name that points to a valid StateMachine")
|
||||
|
||||
|
||||
class ModelWithStateMachine(BaseModelWithStateMachine):
|
||||
StatusChoices = DefaultStatusChoices
|
||||
|
||||
status: models.CharField = BaseModelWithStateMachine.StatusField()
|
||||
retry_at: models.DateTimeField = BaseModelWithStateMachine.RetryAtField()
|
||||
|
||||
state_machine_name: str | None # e.g. 'core.models.ArchiveResultMachine'
|
||||
state_field_name: str = "status"
|
||||
state_machine_attr: str = "sm"
|
||||
bind_events_as_methods: bool = False
|
||||
|
||||
active_state = StatusChoices.STARTED
|
||||
retry_at_field_name: str = "retry_at"
|
||||
|
||||
class Meta(BaseModelWithStateMachine.Meta):
|
||||
abstract = True
|
||||
|
||||
|
||||
class BaseStateMachine(StateMachine):
|
||||
"""
|
||||
Base class for all ArchiveBox state machines.
|
||||
|
||||
Eliminates boilerplate __init__, __repr__, __str__ methods that were
|
||||
duplicated across all 4 state machines (Snapshot, ArchiveResult, Crawl, Binary).
|
||||
|
||||
Subclasses must set model_attr_name to specify the attribute name
|
||||
(e.g., 'snapshot', 'archiveresult', 'crawl', 'binary').
|
||||
|
||||
Example usage:
|
||||
class SnapshotMachine(BaseStateMachine):
|
||||
model_attr_name = 'snapshot'
|
||||
|
||||
# States and transitions...
|
||||
queued = State(value=Snapshot.StatusChoices.QUEUED, initial=True)
|
||||
# ...
|
||||
|
||||
The model instance is accessible via self.{model_attr_name}
|
||||
(e.g., self.snapshot, self.archiveresult, etc.)
|
||||
"""
|
||||
|
||||
model_attr_name: str = "obj" # Override in subclasses
|
||||
|
||||
def __init__(self, obj, *args, **kwargs):
|
||||
setattr(self, self.model_attr_name, obj)
|
||||
super().__init__(obj, *args, **kwargs)
|
||||
|
||||
def _register_callbacks(self, listeners: list[object]):
|
||||
"""Register transition callbacks without scanning the Django model.
|
||||
|
||||
python-statemachine normally treats the wrapped model as a callback
|
||||
listener. That is useful when transition specs point at methods on the
|
||||
domain object, but ArchiveBox keeps all transition guards/actions on the
|
||||
machine classes themselves (`SnapshotMachine.can_start`,
|
||||
`CrawlMachine.enter_sealed`, etc.). Scanning the Django model therefore
|
||||
only adds work: `dir(model)` is large, callback resolution walks that
|
||||
attribute set for every state/transition, and the cost lands on every
|
||||
`.sm` construction.
|
||||
|
||||
Keep support for explicit external listeners, but do not register
|
||||
`self.model` as an implicit listener. If a future machine wants model
|
||||
methods as callbacks, pass that model explicitly as a listener at the
|
||||
callsite so the cost is local and visible.
|
||||
"""
|
||||
self._listeners.update({id(listener): listener for listener in listeners})
|
||||
callbacks = Listeners.from_listeners(
|
||||
(
|
||||
Listener.from_obj(self, skip_attrs=self._protected_attrs),
|
||||
*(Listener.from_obj(listener) for listener in listeners),
|
||||
),
|
||||
)
|
||||
registry = self._callbacks
|
||||
callbacks.resolve(self._specs, registry=registry, allowed_references=SPECS_ALL)
|
||||
|
||||
check_callbacks = self._callbacks.check
|
||||
for visited in iterate_states_and_transitions(self.states):
|
||||
callbacks.resolve(visited._specs, registry=registry, allowed_references=SPECS_ALL)
|
||||
check_callbacks(visited._specs)
|
||||
|
||||
self._callbacks.async_or_sync()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
obj = getattr(self, self.model_attr_name)
|
||||
return f"{self.__class__.__name__}[{obj.id}]"
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.__repr__()
|
||||
def RetryAtField(cls, **kwargs: Any) -> models.DateTimeField:
|
||||
return models.DateTimeField(**{**default_retry_at_field.deconstruct()[3], **kwargs})
|
||||
|
||||
@ -49,8 +49,6 @@ dependencies = [
|
||||
"django-admin-data-views>=0.4.1",
|
||||
"django-object-actions>=4.3.0",
|
||||
"bleach>=6.2.0", # for: stripping unsafe HTML from user-editable titles, notes, labels, tags
|
||||
### State Management
|
||||
"python-statemachine[diagrams]>=2.3.6",
|
||||
### CLI / Logging
|
||||
"click>=8.3.1", # for: nicer CLI command + argument definitions
|
||||
"rich>=14.2.0", # for: pretty CLI output
|
||||
|
||||
34
uv.lock
34
uv.lock
@ -149,7 +149,6 @@ dependencies = [
|
||||
{ name = "py-machineid", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
{ name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
{ name = "pydantic-settings", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
{ name = "python-statemachine", extra = ["diagrams"], marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
{ name = "requests", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
{ name = "rich", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
{ name = "rich-click", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
@ -253,7 +252,6 @@ requires-dist = [
|
||||
{ name = "pydantic", specifier = ">=2.8.0" },
|
||||
{ name = "pydantic-settings", specifier = ">=2.5.2" },
|
||||
{ name = "python-ldap", marker = "extra == 'ldap'", specifier = ">=3.4.3" },
|
||||
{ name = "python-statemachine", extras = ["diagrams"], specifier = ">=2.3.6" },
|
||||
{ name = "requests", specifier = ">=2.32.3" },
|
||||
{ name = "requests-tracker", marker = "extra == 'debug'", specifier = ">=0.3.3" },
|
||||
{ name = "rich", specifier = ">=14.2.0" },
|
||||
@ -1798,17 +1796,6 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/77/c1/6e422f34e569cf8e18df68d1939c81c099d2b61e4f7d9621c8a77560799c/pydantic_settings-2.14.2-py3-none-any.whl", hash = "sha256:a20c97b37910b6550d5ea50fbcc2d4187defe58cd57070b73863d069419c9440", size = 61715, upload-time = "2026-06-19T13:44:55.02Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pydot"
|
||||
version = "4.0.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pyparsing", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/50/35/b17cb89ff865484c6a20ef46bf9d95a5f07328292578de0b295f4a6beec2/pydot-4.0.1.tar.gz", hash = "sha256:c2148f681c4a33e08bf0e26a9e5f8e4099a82e0e2a068098f32ce86577364ad5", size = 162594, upload-time = "2025-06-17T20:09:56.454Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/7e/32/a7125fb28c4261a627f999d5fb4afff25b523800faed2c30979949d6facd/pydot-4.0.1-py3-none-any.whl", hash = "sha256:869c0efadd2708c0be1f916eb669f3d664ca684bc57ffb7ecc08e70d5e93fee6", size = 37087, upload-time = "2025-06-17T20:09:55.25Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pygments"
|
||||
@ -1831,14 +1818,6 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/54/18/1dd71c9b43192ab83f1d531ad6002dc81108ac36c475f79fb7a295abe2f4/pyopenssl-26.3.0-py3-none-any.whl", hash = "sha256:46367f8f66b92271e6d218da9c87607e1ef5a0bc5c8dea5bb3db82f395c385a3", size = 56008, upload-time = "2026-06-12T20:28:05.999Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyparsing"
|
||||
version = "3.3.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f3/91/9c6ee907786a473bf81c5f53cf703ba0957b23ab84c264080fb5a450416f/pyparsing-3.3.2.tar.gz", hash = "sha256:c777f4d763f140633dcb6d8a3eda953bf7a214dc4eff598413c070bcdc117cbc", size = 6851574, upload-time = "2026-01-21T03:57:59.36Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyright"
|
||||
@ -1974,19 +1953,6 @@ dependencies = [
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b2/f4/60edeb794bbc9ed0ff2149bbaeec605f3ed331766459d195832ecbd0ba2d/python_ldap-3.4.7.tar.gz", hash = "sha256:bacd9fb680d20263d8570ade1cf234d90d281149a8beb4f079dd8f33f7613dc8", size = 387477, upload-time = "2026-05-20T13:41:04.358Z" }
|
||||
|
||||
[[package]]
|
||||
name = "python-statemachine"
|
||||
version = "3.2.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/a3/88/d24d2475069c96bbd752414863d8ffc26ef25f5179c87fda5869ba8cb5d5/python_statemachine-3.2.0.tar.gz", hash = "sha256:44b98cb9bb1081891ef6efa907c821d05cae5314a2b99713e56d56a9c962de08", size = 745821, upload-time = "2026-06-17T02:25:41.347Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d8/e3/25d1076e47554b559332bba62fb4abdc01704ea960f773c8db21b4ea1952/python_statemachine-3.2.0-py3-none-any.whl", hash = "sha256:8909916fc21208680f737d33c891a5331a1a0eeeaf3fe40efa278cb056855aee", size = 152809, upload-time = "2026-06-17T02:25:39.766Z" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
diagrams = [
|
||||
{ name = "pydot", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytz"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user