more StateMachine, Actor, and Orchestrator improvements

archivebox/actors/actor.py

@@ -2,78 +2,240 @@ __package__ = 'archivebox.actors'
 
 import os
 import time
-from abc import ABC, abstractmethod
+from typing import ClassVar, Generic, TypeVar, Any, Literal, Type, Iterable, cast, get_args
-from typing import ClassVar, Generic, TypeVar, Any, cast, Literal, Type
+from datetime import timedelta
-from django.utils.functional import classproperty
-
-from rich import print
-import psutil
-
-from django import db
-from django.db import models
-from django.db.models import QuerySet
 from multiprocessing import Process, cpu_count
 from threading import Thread, get_native_id
 
+import psutil
+from rich import print
+from statemachine import State, StateMachine, registry
+from statemachine.mixins import MachineMixin
+
+from django import db
+from django.db.models import QuerySet, sql, Q
+from django.db.models import Model as DjangoModel
+from django.utils import timezone
+from django.utils.functional import classproperty
+
+from .models import ModelWithStateMachine
+
 # from archivebox.logging_util import TimedProgress
 
+class ActorObjectAlreadyClaimed(Exception):
+    """Raised when the Actor tries to claim the next object from the queue but it's already been claimed by another Actor"""
+    pass
+
+class ActorQueueIsEmpty(Exception):
+    """Raised when the Actor tries to get the next object from the queue but it's empty"""
+    pass
+
+CPU_COUNT = cpu_count()
+DEFAULT_MAX_TICK_TIME = 60
+DEFAULT_MAX_CONCURRENT_ACTORS = min(max(2, int(CPU_COUNT * 0.6)), 8)   # 2 < 60% * num available cpu cores < 8
+
+limit = lambda n, max: min(n, max)
+
 LaunchKwargs = dict[str, Any]
+ObjectState = State | str
+ObjectStateList = Iterable[ObjectState]
 
-ModelType = TypeVar('ModelType', bound=models.Model)
+ModelType = TypeVar('ModelType', bound=ModelWithStateMachine)
 
-class ActorType(ABC, Generic[ModelType]):
+class ActorType(Generic[ModelType]):
     """
     Base class for all actors. Usage:
-    class FaviconActor(ActorType[ArchiveResult]):
-        QUERYSET: ClassVar[QuerySet] = ArchiveResult.objects.filter(status='queued', extractor='favicon')
-        CLAIM_WHERE: ClassVar[str] = 'status = "queued" AND extractor = "favicon"'
-        CLAIM_ORDER: ClassVar[str] = 'created_at DESC'
-        ATOMIC: ClassVar[bool] = True
     
-        def claim_sql_set(self, obj: ArchiveResult) -> str:
+    class FaviconActor(ActorType[FaviconArchiveResult]):
-            # SQL fields to update atomically while claiming an object from the queue
+        FINAL_STATES: ClassVar[tuple[str, ...]] = ('succeeded', 'failed', 'skipped')
-            retry_at = datetime.now() + timedelta(seconds=self.MAX_TICK_TIME)
+        ACTIVE_STATE: ClassVar[str] = 'started'
-            return f"status = 'started', locked_by = {self.pid}, retry_at = {retry_at}"
         
-        def tick(self, obj: ArchiveResult) -> None:
+        @classmethod
-            run_favicon_extractor(obj)
+        def qs(cls) -> QuerySet[FaviconArchiveResult]:
-            ArchiveResult.objects.filter(pk=obj.pk, status='started').update(status='success')
+            return ArchiveResult.objects.filter(extractor='favicon')   # or leave the default: FaviconArchiveResult.objects.all()
     """
+    
+    ### Class attributes (defined on the class at compile-time when ActorType[MyModel] is defined)
+    Model: Type[ModelType]
+    StateMachineClass: Type[StateMachine]
+    
+    STATE_FIELD_NAME: ClassVar[str]
+    ACTIVE_STATE: ClassVar[ObjectState]
+    FINAL_STATES: ClassVar[ObjectStateList]
+    EVENT_NAME: ClassVar[str] = 'tick'                                    # the event name to trigger on the obj.sm: StateMachine (usually 'tick')
+    
+    CLAIM_ORDER: ClassVar[tuple[str, ...]] = ('retry_at',)                # the .order(*args) to claim the queue objects in, use ('?',) for random order
+    CLAIM_FROM_TOP_N: ClassVar[int] = CPU_COUNT * 10                      # the number of objects to consider when atomically getting the next object from the queue
+    CLAIM_ATOMIC: ClassVar[bool] = True                                   # whether to atomically fetch+claim the next object in one query, or fetch and lock it in two queries
+    
+    MAX_TICK_TIME: ClassVar[int] = DEFAULT_MAX_TICK_TIME                  # maximum duration in seconds to process a single object
+    MAX_CONCURRENT_ACTORS: ClassVar[int] = DEFAULT_MAX_CONCURRENT_ACTORS  # maximum number of concurrent actors that can be running at once
+    
+    _SPAWNED_ACTOR_PIDS: ClassVar[list[psutil.Process]] = []      # used to record all the pids of Actors spawned on the class
+    
+    ### Instance attributes (only used within an actor instance inside a spawned actor thread/process)
     pid: int
     idle_count: int = 0
     launch_kwargs: LaunchKwargs = {}
     mode: Literal['thread', 'process'] = 'process'
     
-    MAX_CONCURRENT_ACTORS: ClassVar[int] = min(max(2, int(cpu_count() * 0.6)), 8)   # min 2, max 8, up to 60% of available cpu cores
+    def __init_subclass__(cls) -> None:
-    MAX_TICK_TIME: ClassVar[int] = 60                          # maximum duration in seconds to process a single object
+        """
-    
+        Executed at class definition time (i.e. during import of any file containing class MyActor(ActorType[MyModel]): ...).
-    QUERYSET: ClassVar[QuerySet]                      # the QuerySet to claim objects from
+        Loads the django Model from the Generic[ModelType] TypeVar arg and populates any missing class-level config using it.
-    CLAIM_WHERE: ClassVar[str] = 'status = "queued"'  # the WHERE clause to filter the objects when atomically getting the next object from the queue
+        """
-    CLAIM_SET: ClassVar[str] = 'status = "started"'   # the SET clause to claim the object when atomically getting the next object from the queue
+        if getattr(cls, 'Model', None) is None:
-    CLAIM_ORDER: ClassVar[str] = 'created_at DESC'    # the ORDER BY clause to sort the objects with when atomically getting the next object from the queue
+            cls.Model = cls._get_model_from_generic_typevar()
-    CLAIM_FROM_TOP: ClassVar[int] = MAX_CONCURRENT_ACTORS * 10  # the number of objects to consider when atomically getting the next object from the queue
+        cls._populate_missing_classvars_from_model(cls.Model)
-    ATOMIC: ClassVar[bool] = True                     # whether to atomically fetch+claim the nextobject in one step, or fetch and lock it in two steps
-    
-    # model_type: Type[ModelType]
-    
-    _SPAWNED_ACTOR_PIDS: ClassVar[list[psutil.Process]] = []   # record all the pids of Actors spawned by this class
     
     def __init__(self, mode: Literal['thread', 'process']|None=None, **launch_kwargs: LaunchKwargs):
+        """
+        Executed right before the Actor is spawned to create a unique Actor instance for that thread/process.
+        actor_instance.runloop() is then executed from inside the newly spawned thread/process.
+        """
         self.mode = mode or self.mode
         self.launch_kwargs = launch_kwargs or dict(self.launch_kwargs)
     
+
+    ### Private Helper Methods: Not desiged to be overridden by subclasses or called by anything outside of this class
+    
     @classproperty
     def name(cls) -> str:
         return cls.__name__  # type: ignore
     
     def __str__(self) -> str:
-        return self.__repr__()
+        return repr(self)
     
     def __repr__(self) -> str:
-        """FaviconActor[pid=1234]"""
+        """-> FaviconActor[pid=1234]"""
         label = 'pid' if self.mode == 'process' else 'tid'
         return f'[underline]{self.name}[/underline]\\[{label}={self.pid}]'
     
+    @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)
+    
+    @staticmethod
+    def _sql_for_select_top_n_candidates(qs: QuerySet, claim_from_top_n: int=CLAIM_FROM_TOP_N) -> tuple[str, tuple[Any, ...]]:
+        """Get the SQL for selecting the top N candidates from the queue (to claim one from)"""
+        queryset = qs.only('id')[:claim_from_top_n]
+        select_sql, select_params = compile_sql_select(queryset)
+        return select_sql, select_params
+    
+    @staticmethod
+    def _sql_for_update_claimed_obj(qs: QuerySet, update_kwargs: dict[str, Any]) -> tuple[str, tuple[Any, ...]]:
+        """Get the SQL for updating a claimed object to mark it as ACTIVE"""
+        # qs.update(status='started', retry_at=<now + MAX_TICK_TIME>)
+        update_sql, update_params = compile_sql_update(qs, update_kwargs=update_kwargs)
+        # e.g. UPDATE core_archiveresult SET status='%s', retry_at='%s' WHERE status NOT IN ('succeeded', 'failed', 'sealed', 'started') AND retry_at <= '2024-11-04 10:14:33.240903'
+        return update_sql, update_params
+    
+    @classmethod
+    def _get_model_from_generic_typevar(cls) -> Type[ModelType]:
+        """Get the django Model from the Generic[ModelType] TypeVar arg (and check that it inherits from django.db.models.Model)"""
+        # cls.__orig_bases__ is non-standard and may be removed in the future! if this breaks,
+        # we can just require the inerited class to define the Model as a classvar manually, e.g.:
+        #     class SnapshotActor(ActorType[Snapshot]):
+        #         Model: ClassVar[Type[Snapshot]] = Snapshot
+        # https://stackoverflow.com/questions/57706180/generict-base-class-how-to-get-type-of-t-from-within-instance
+        Model = get_args(cls.__orig_bases__[0])[0]   # type: ignore
+        assert issubclass(Model, DjangoModel), f'{cls.__name__}.Model must be a valid django Model'
+        return cast(Type[ModelType], Model)
+    
+    @staticmethod
+    def _get_state_machine_cls(Model: Type[ModelType]) -> Type[StateMachine]:
+        """Get the StateMachine class for the given django Model that inherits from MachineMixin"""
+        assert issubclass(Model, MachineMixin), f'{Model.__name__} must inherit from MachineMixin and define a .state_machine_name: str'
+        model_state_machine_name = getattr(Model, 'state_machine_name', None)
+        if model_state_machine_name:
+            StateMachine = registry.get_machine_cls(model_state_machine_name)
+            assert issubclass(StateMachine, StateMachine)
+            return StateMachine
+        raise NotImplementedError(f'ActorType[{Model.__name__}] must define .state_machine_name: str that points to a valid StateMachine')
+    
+    @classmethod
+    def _get_state_machine_instance(cls, obj: ModelType) -> StateMachine:
+        """Get the StateMachine instance for the given django Model instance (and check that it is a valid instance of cls.StateMachineClass)"""
+        obj_statemachine = None
+        state_machine_attr = getattr(obj, 'state_machine_attr', 'sm')
+        try:
+            obj_statemachine = getattr(obj, state_machine_attr)
+        except Exception:
+            pass
+        
+        if not isinstance(obj_statemachine, cls.StateMachineClass):
+            raise Exception(f'{cls.__name__}: Failed to find a valid StateMachine instance at {type(obj).__name__}.{state_machine_attr}')
+            
+        return obj_statemachine
+    
+    @classmethod
+    def _populate_missing_classvars_from_model(cls, Model: Type[ModelType]):
+        """Check that the class variables are set correctly based on the ModelType"""
+        
+        # check that Model is the same as the Generic[ModelType] parameter in the class definition
+        cls.Model = getattr(cls, 'Model', None) or Model
+        if cls.Model != Model:
+            raise ValueError(f'{cls.__name__}.Model must be set to the same Model as the Generic[ModelType] parameter in the class definition')
+        
+        # check that Model has a valid StateMachine with the required event defined on it
+        cls.StateMachineClass = getattr(cls, 'StateMachineClass', None) or cls._get_state_machine_cls(cls.Model)
+        assert isinstance(cls.EVENT_NAME, str), f'{cls.__name__}.EVENT_NAME must be a str, got: {type(cls.EVENT_NAME).__name__} instead'
+        assert hasattr(cls.StateMachineClass, cls.EVENT_NAME), f'StateMachine {cls.StateMachineClass.__name__} must define a {cls.EVENT_NAME} event ({cls.__name__}.EVENT_NAME = {cls.EVENT_NAME})'
+        
+        # check that Model uses .id as its primary key field
+        primary_key_field = cls.Model._meta.pk.name
+        if primary_key_field != 'id':
+            raise NotImplementedError(f'Actors currently only support models that use .id as their primary key field ({cls.__name__} uses {cls.__name__}.{primary_key_field} as primary key)')
+        
+        # check for STATE_FIELD_NAME classvar or set it from the model's state_field_name attr
+        if not getattr(cls, 'STATE_FIELD_NAME', None):
+            if hasattr(cls.Model, 'state_field_name'):
+                cls.STATE_FIELD_NAME = getattr(cls.Model, 'state_field_name')
+            else:
+                raise NotImplementedError(f'{cls.__name__} must define a STATE_FIELD_NAME: ClassVar[str] (e.g. "status") or have a .state_field_name attr on its Model')
+        assert isinstance(cls.STATE_FIELD_NAME, str), f'{cls.__name__}.STATE_FIELD_NAME must be a str, got: {type(cls.STATE_FIELD_NAME).__name__} instead'
+            
+        # check for FINAL_STATES classvar or set it from the model's final_states attr
+        if not getattr(cls, 'FINAL_STATES', None):
+            cls.FINAL_STATES = cls.StateMachineClass.final_states
+            if not cls.FINAL_STATES:
+                raise NotImplementedError(f'{cls.__name__} must define a non-empty FINAL_STATES: ClassVar[list[str]] (e.g. ["sealed"]) or have a {cls.Model.__name__}.state_machine_name pointing to a StateMachine that provides .final_states')
+        cls.FINAL_STATES = [cls._state_to_str(state) for state in cls.FINAL_STATES]
+        assert all(isinstance(state, str) for state in cls.FINAL_STATES), f'{cls.__name__}.FINAL_STATES must be a list[str], got: {type(cls.FINAL_STATES).__name__} instead'
+        
+        # check for ACTIVE_STATE classvar or set it from the model's active_state attr
+        if not getattr(cls, 'ACTIVE_STATE', None):
+            raise NotImplementedError(f'{cls.__name__} must define an ACTIVE_STATE: ClassVar[State] (e.g. SnapshotMachine.started) ({cls.Model.__name__}.{cls.STATE_FIELD_NAME} gets set to this value to mark objects as actively processing)')
+        assert isinstance(cls.ACTIVE_STATE, (State, str)), f'{cls.__name__}.ACTIVE_STATE must be a statemachine.State | str, got: {type(cls.ACTIVE_STATE).__name__} instead'
+        
+        # check the other ClassVar attributes for valid values
+        assert cls.CLAIM_ORDER and isinstance(cls.CLAIM_ORDER, tuple) and all(isinstance(order, str) for order in cls.CLAIM_ORDER), f'{cls.__name__}.CLAIM_ORDER must be a non-empty tuple[str, ...], got: {type(cls.CLAIM_ORDER).__name__} instead'
+        assert cls.CLAIM_FROM_TOP_N > 0, f'{cls.__name__}.CLAIM_FROM_TOP_N must be a positive int, got: {cls.CLAIM_FROM_TOP_N} instead'
+        assert cls.MAX_TICK_TIME >= 1, f'{cls.__name__}.MAX_TICK_TIME must be a positive int > 1, got: {cls.MAX_TICK_TIME} instead'
+        assert cls.MAX_CONCURRENT_ACTORS >= 1, f'{cls.__name__}.MAX_CONCURRENT_ACTORS must be a positive int >=1, got: {cls.MAX_CONCURRENT_ACTORS} instead'
+        assert isinstance(cls.CLAIM_ATOMIC, bool), f'{cls.__name__}.CLAIM_ATOMIC must be a bool, got: {cls.CLAIM_ATOMIC} instead'
+
+    @classmethod
+    def _fork_actor_as_thread(cls, **launch_kwargs: LaunchKwargs) -> int:
+        """Spawn a new background thread running the actor's runloop"""
+        actor = cls(mode='thread', **launch_kwargs)
+        bg_actor_thread = Thread(target=actor.runloop)
+        bg_actor_thread.start()
+        assert bg_actor_thread.native_id is not None
+        return bg_actor_thread.native_id
+    
+    @classmethod
+    def _fork_actor_as_process(cls, **launch_kwargs: LaunchKwargs) -> int:
+        """Spawn a new background process running the actor's runloop"""
+        actor = cls(mode='process', **launch_kwargs)
+        bg_actor_process = Process(target=actor.runloop)
+        bg_actor_process.start()
+        assert bg_actor_process.pid is not None
+        cls._SPAWNED_ACTOR_PIDS.append(psutil.Process(pid=bg_actor_process.pid))
+        return bg_actor_process.pid
+    
+    
     ### Class Methods: Called by Orchestrator on ActorType class before it has been spawned
     
     @classmethod
@@ -94,71 +256,92 @@ class ActorType(ABC, Generic[ModelType]):
         if not queue_length:                                      # queue is empty, spawn 0 actors
             return []
         
-        actors_to_spawn: list[LaunchKwargs] = []
+        # WARNING:
-        max_spawnable = cls.MAX_CONCURRENT_ACTORS - len(running_actors)
+        # spawning new actors processes is slow/expensive, avoid spawning many actors at once in a single orchestrator tick.
+        # limit to spawning 1 or 2 at a time per orchestrator tick, and let the next tick handle starting another couple.
+        # DONT DO THIS:
+        # if queue_length > 20:                      # queue is extremely long, spawn maximum actors at once!
+        #   num_to_spawn_this_tick = cls.MAX_CONCURRENT_ACTORS
         
-        # spawning new actors is expensive, avoid spawning all the actors at once. To stagger them,
+        if queue_length > 10:    
-        # let the next orchestrator tick handle starting another 2 on the next tick()
+            num_to_spawn_this_tick = 2  # spawn more actors per tick if queue is long
-        # if queue_length > 10:                                   # queue is long, spawn as many as possible
+        else:
-        #   actors_to_spawn += max_spawnable * [{}]
+            num_to_spawn_this_tick = 1  # spawn fewer actors per tick if queue is short
         
-        if queue_length > 4:                                    # queue is medium, spawn 1 or 2 actors
+        num_remaining = cls.MAX_CONCURRENT_ACTORS - len(running_actors)
-            actors_to_spawn += min(2, max_spawnable) * [{**cls.launch_kwargs}]
+        num_to_spawn_now: int = limit(num_to_spawn_this_tick, num_remaining)
-        else:                                                     # queue is short, spawn 1 actor
+        
-            actors_to_spawn += min(1, max_spawnable) * [{**cls.launch_kwargs}]
+        actors_launch_kwargs: list[LaunchKwargs] = num_to_spawn_now * [{**cls.launch_kwargs}]
-        return actors_to_spawn
+        return actors_launch_kwargs
         
     @classmethod
     def start(cls, mode: Literal['thread', 'process']='process', **launch_kwargs: LaunchKwargs) -> int:
         if mode == 'thread':
-            return cls.fork_actor_as_thread(**launch_kwargs)
+            return cls._fork_actor_as_thread(**launch_kwargs)
         elif mode == 'process':
-            return cls.fork_actor_as_process(**launch_kwargs)
+            return cls._fork_actor_as_process(**launch_kwargs)
         raise ValueError(f'Invalid actor mode: {mode} must be "thread" or "process"')
     
-    @classmethod
+    @classproperty
-    def fork_actor_as_thread(cls, **launch_kwargs: LaunchKwargs) -> int:
+    def qs(cls) -> QuerySet[ModelType]:
-        """Spawn a new background thread running the actor's runloop"""
+        """
-        actor = cls(mode='thread', **launch_kwargs)
+        Get the unfiltered and unsorted QuerySet of all objects that this Actor might care about.
-        bg_actor_thread = Thread(target=actor.runloop)
+        Override this in the subclass to define the QuerySet of objects that the Actor is going to poll for new work.
-        bg_actor_thread.start()
+        (don't limit, order, or filter this by retry_at or status yet, Actor.get_queue() handles that part)
-        assert bg_actor_thread.native_id is not None
+        """
-        return bg_actor_thread.native_id
+        return cls.Model.objects.all()
+    
+    @classproperty
+    def final_q(cls) -> Q:
+        """Get the filter for objects that are in a final state"""
+        return Q(**{f'{cls.STATE_FIELD_NAME}__in': [cls._state_to_str(s) for s in cls.FINAL_STATES]})
+    
+    @classproperty
+    def active_q(cls) -> Q:
+        """Get the filter for objects that are actively processing right now"""
+        return Q(**{cls.STATE_FIELD_NAME: cls._state_to_str(cls.ACTIVE_STATE)})   # e.g. Q(status='started')
+    
+    @classproperty
+    def stalled_q(cls) -> Q:
+        """Get the filter for objects that are marked active but have timed out"""
+        return cls.active_q & Q(retry_at__lte=timezone.now())                     # e.g. Q(status='started') AND Q(<retry_at is in the past>)
+    
+    @classproperty
+    def future_q(cls) -> Q:
+        """Get the filter for objects that have a retry_at in the future"""
+        return Q(retry_at__gt=timezone.now())
+    
+    @classproperty
+    def pending_q(cls) -> Q:
+        """Get the filter for objects that are ready for processing."""
+        return ~(cls.active_q) & ~(cls.final_q) & ~(cls.future_q)
     
     @classmethod
-    def fork_actor_as_process(cls, **launch_kwargs: LaunchKwargs) -> int:
+    def get_queue(cls, sort: bool=True) -> QuerySet[ModelType]:
-        """Spawn a new background process running the actor's runloop"""
+        """
-        actor = cls(mode='process', **launch_kwargs)
+        Get the sorted and filtered QuerySet of objects that are ready for processing.
-        bg_actor_process = Process(target=actor.runloop)
+        e.g. qs.exclude(status__in=('sealed', 'started'), retry_at__gt=timezone.now()).order_by('retry_at')
-        bg_actor_process.start()
+        """
-        assert bg_actor_process.pid is not None
+        unsorted_qs = cls.qs.filter(cls.pending_q)
-        cls._SPAWNED_ACTOR_PIDS.append(psutil.Process(pid=bg_actor_process.pid))
+        return unsorted_qs.order_by(*cls.CLAIM_ORDER) if sort else unsorted_qs
-        return bg_actor_process.pid
 
-    @classmethod
+    ### Instance Methods: Only called from within Actor instance after it has been spawned (i.e. forked as a thread or process)
-    def get_model(cls) -> Type[ModelType]:
-        # wish this was a @classproperty but Generic[ModelType] return type cant be statically inferred for @classproperty
-        return cls.QUERYSET.model
-    
-    @classmethod
-    def get_queue(cls) -> QuerySet:
-        """override this to provide your queryset as the queue"""
-        # return ArchiveResult.objects.filter(status='queued', extractor__in=('pdf', 'dom', 'screenshot'))
-        return cls.QUERYSET
-    
-    ### Instance Methods: Called by Actor after it has been spawned (i.e. forked as a thread or process)
     
     def runloop(self):
         """The main runloop that starts running when the actor is spawned (as subprocess or thread) and exits when the queue is empty"""
         self.on_startup()
+        obj_to_process: ModelType | None = None
+        last_error: BaseException | None = None
         try:
             while True:
-                obj_to_process: ModelType | None = None
+                # Get the next object to process from the queue
                 try:
                     obj_to_process = cast(ModelType, self.get_next(atomic=self.atomic))
-                except Exception:
+                except (ActorQueueIsEmpty, ActorObjectAlreadyClaimed) as err:
-                    pass
+                    last_error = err
+                    obj_to_process = None
                 
+                # Handle the case where there is no next object to process
                 if obj_to_process:
                     self.idle_count = 0   # reset idle count if we got an object
                 else:
@@ -170,118 +353,126 @@ class ActorType(ABC, Generic[ModelType]):
                         time.sleep(1)
                         continue
                 
+                # Process the object by triggering its StateMachine.tick() method
                 self.on_tick_start(obj_to_process)
-                
-                # Process the object
                 try:
                     self.tick(obj_to_process)
                 except Exception as err:
-                    print(f'[red]🏃‍♂️ ERROR: {self}.tick()[/red]', err)
+                    last_error = err
+                    # print(f'[red]🏃‍♂️ {self}.tick()[/red] {obj_to_process} ERROR: [red]{type(err).__name__}: {err}[/red]')
                     db.connections.close_all()                         # always reset the db connection after an exception to clear any pending transactions
                     self.on_tick_exception(obj_to_process, err)
                 finally:
                     self.on_tick_end(obj_to_process)
 
-            self.on_shutdown(err=None)
         except BaseException as err:
+            last_error = err
             if isinstance(err, KeyboardInterrupt):
                 print()
             else:
-                print(f'\n[red]🏃‍♂️ {self}.runloop() FATAL:[/red]', err.__class__.__name__, err)
+                print(f'\n[red]🏃‍♂️ {self}.runloop() FATAL:[/red] {type(err).__name__}: {err}')
-            self.on_shutdown(err=err)
+                print(f'    Last processed object: {obj_to_process}')
+                raise
+        finally:
+            self.on_shutdown(last_obj=obj_to_process, last_error=last_error)
+    
+    def get_update_kwargs_to_claim_obj(self) -> dict[str, Any]:
+        """
+        Get the field values needed to mark an pending obj_to_process as being actively processing (aka claimed)
+        by the current Actor. returned kwargs will be applied using: qs.filter(id=obj_to_process.id).update(**kwargs).
+        F() expressions are allowed in field values if you need to update a field based on its current value.
+        Can be a defined as a normal method (instead of classmethod) on subclasses if it needs to access instance vars.
+        """
+        return {
+            self.STATE_FIELD_NAME: self.ACTIVE_STATE,
+            'retry_at': timezone.now() + timedelta(seconds=self.MAX_TICK_TIME),
+        }
     
     def get_next(self, atomic: bool | None=None) -> ModelType | None:
         """get the next object from the queue, atomically locking it if self.atomic=True"""
-        if atomic is None:
+        atomic = self.CLAIM_ATOMIC if atomic is None else atomic
-            atomic = self.ATOMIC
-
         if atomic:
             # fetch and claim the next object from in the queue in one go atomically
             obj = self.get_next_atomic()
         else:
             # two-step claim: fetch the next object and lock it in a separate query
-            obj = self.get_queue().last()
+            obj = self.get_next_non_atomic()
-            assert obj and self.lock_next(obj), f'Unable to fetch+lock the next {self.get_model().__name__} ojbect from {self}.QUEUE'
         return obj
     
-    def lock_next(self, obj: ModelType) -> bool:
+    def get_next_non_atomic(self) -> ModelType:
-        """override this to implement a custom two-step (non-atomic)lock mechanism"""
-        # For example:
-        # assert obj._model.objects.filter(pk=obj.pk, status='queued').update(status='started', locked_by=self.pid)
-        # Not needed if using get_next_and_lock() to claim the object atomically
-        # print(f'[blue]🏃‍♂️ {self}.lock()[/blue]', obj.abid or obj.id)
-        return True
-    
-    def claim_sql_where(self) -> str:
-        """override this to implement a custom WHERE clause for the atomic claim step e.g. "status = 'queued' AND locked_by = NULL" """
-        return self.CLAIM_WHERE
-    
-    def claim_sql_set(self) -> str:
-        """override this to implement a custom SET clause for the atomic claim step e.g. "status = 'started' AND locked_by = {self.pid}" """
-        return self.CLAIM_SET
-    
-    def claim_sql_order(self) -> str:
-        """override this to implement a custom ORDER BY clause for the atomic claim step e.g. "created_at DESC" """
-        return self.CLAIM_ORDER
-    
-    def claim_from_top(self) -> int:
-        """override this to implement a custom number of objects to consider when atomically claiming the next object from the top of the queue"""
-        return self.CLAIM_FROM_TOP
-        
-    def get_next_atomic(self, shallow: bool=True) -> ModelType | None:
         """
-        claim a random object from the top n=50 objects in the queue (atomically updates status=queued->started for claimed object)
+        Naiively selects the top/first object from self.get_queue().order_by(*self.CLAIM_ORDER),
-        optimized for minimizing contention on the queue with other actors selecting from the same list
+        then claims it by running .update(status='started', retry_at=<now + MAX_TICK_TIME>).
-        slightly faster than claim_any_obj() which selects randomly from the entire queue but needs to know the total count
+        
+        Do not use this method if there is more than one Actor racing to get objects from the same queue,
+        it will be slow/buggy as they'll compete to lock the same object at the same time (TOCTTOU race).
         """
-        Model = self.get_model()                                     # e.g. ArchiveResult
+        obj = self.get_queue().first()
-        table = f'{Model._meta.app_label}_{Model._meta.model_name}'  # e.g. core_archiveresult
+        if obj is None:
+            raise ActorQueueIsEmpty(f'No next object available in {self}.get_queue()')
         
-        where_sql = self.claim_sql_where()
+        locked = self.get_queue().filter(id=obj.id).update(**self.get_update_kwargs_to_claim_obj())
-        set_sql = self.claim_sql_set()
+        if not locked:
-        order_by_sql = self.claim_sql_order()
+            raise ActorObjectAlreadyClaimed(f'Unable to lock the next {self.Model.__name__} object from {self}.get_queue().first()')
-        choose_from_top = self.claim_from_top()
+        return obj
         
-        with db.connection.cursor() as cursor:
+    def get_next_atomic(self) -> ModelType | None:
-            # subquery gets the pool of the top 50 candidates sorted by sort and order
+        """
-            # main query selects a random one from that pool
+        Selects the top n=50 objects from the queue and atomically claims a random one from that set.
-            cursor.execute(f"""
+        This approach safely minimizes contention with other Actors trying to select from the same Queue.
-                UPDATE {table} 
+
-                SET {set_sql}
+        The atomic query is roughly equivalent to the following:  (all done in one SQL query to avoid a TOCTTOU race)
-                WHERE {where_sql} and id = (
+            top_candidates are selected from:   qs.order_by(*CLAIM_ORDER).only('id')[:CLAIM_FROM_TOP_N]
-                    SELECT id FROM (
+            a single candidate is chosen using: qs.filter(id__in=top_n_candidates).order_by('?').first()
-                        SELECT id FROM {table}
+            the chosen obj is claimed using:    qs.filter(id=chosen_obj).update(status=ACTIVE_STATE, retry_at=<now + MAX_TICK_TIME>)
-                        WHERE {where_sql}
+        """
-                        ORDER BY {order_by_sql}
+        # TODO: if we switch from SQLite to PostgreSQL in the future, we should change this
-                        LIMIT {choose_from_top}
+        # to use SELECT FOR UPDATE instead of a subquery + ORDER BY RANDOM() LIMIT 1
+        
+        # e.g. SELECT id FROM core_archiveresult WHERE status NOT IN (...) AND retry_at <= '...' ORDER BY retry_at ASC LIMIT 50
+        qs = self.get_queue()
+        select_top_canidates_sql, select_params = self._sql_for_select_top_n_candidates(qs=qs)
+        assert select_top_canidates_sql.startswith('SELECT ')
+        
+        # e.g. UPDATE core_archiveresult SET status='%s', retry_at='%s' WHERE status NOT IN (...) AND retry_at <= '...'
+        update_claimed_obj_sql, update_params = self._sql_for_update_claimed_obj(qs=qs, update_kwargs=self.get_update_kwargs_to_claim_obj())
+        assert update_claimed_obj_sql.startswith('UPDATE ')
+        db_table = self.Model._meta.db_table  # e.g. core_archiveresult
+        
+        # subquery gets the pool of the top candidates e.g. self.get_queue().only('id')[:CLAIM_FROM_TOP_N]
+        # main query selects a random one from that pool, and claims it using .update(status=ACTIVE_STATE, retry_at=<now + MAX_TICK_TIME>)
+        # this is all done in one atomic SQL query to avoid TOCTTOU race conditions (as much as possible)
+        atomic_select_and_update_sql = f"""
+            {update_claimed_obj_sql} AND "{db_table}"."id" = (
+                SELECT "{db_table}"."id" FROM (
+                    {select_top_canidates_sql}
                 ) candidates
                 ORDER BY RANDOM()
                 LIMIT 1
             )
-                RETURNING id;
+            RETURNING *;
-            """)
+        """
-            result = cursor.fetchone()
+        try:
+            return self.Model.objects.raw(atomic_select_and_update_sql, (*update_params, *select_params))[0]
+        except KeyError:
+            if self.get_queue().exists():
+                raise ActorObjectAlreadyClaimed(f'Unable to lock the next {self.Model.__name__} object from {self}.get_queue().first()')
+            else:
+                raise ActorQueueIsEmpty(f'No next object available in {self}.get_queue()')
 
-            if result is None:
+    def tick(self, obj_to_process: ModelType) -> None:
-                return None           # If no rows were claimed, return None
+        """Call the object.sm.tick() method to process the object"""
+        print(f'[blue]🏃‍♂️ {self}.tick()[/blue] {obj_to_process}')
         
-            if shallow:
+        # get the StateMachine instance from the object
-                # shallow: faster, returns potentially incomplete object instance missing some django auto-populated fields:
+        obj_statemachine = self._get_state_machine_instance(obj_to_process)
-                columns = [col[0] for col in cursor.description or ['id']]
-                return Model(**dict(zip(columns, result)))
         
-            # if not shallow do one extra query to get a more complete object instance (load it fully from scratch)
+        # trigger the event on the StateMachine instance
-            return Model.objects.get(id=result[0])
+        obj_tick_method = getattr(obj_statemachine, self.EVENT_NAME)  # e.g. obj_statemachine.tick()
+        obj_tick_method()
         
-    @abstractmethod
+        # save the object to persist any state changes
-    def tick(self, obj: ModelType) -> None:
+        obj_to_process.save()
-        """override this to process the object"""
-        print(f'[blue]🏃‍♂️ {self}.tick()[/blue]', obj.abid or obj.id)
-        # For example:
-        # do_some_task(obj)
-        # do_something_else(obj)
-        # obj._model.objects.filter(pk=obj.pk, status='started').update(status='success')
-        raise NotImplementedError('tick() must be implemented by the Actor subclass')
         
     def on_startup(self) -> None:
         if self.mode == 'thread':
@@ -290,24 +481,91 @@ class ActorType(ABC, Generic[ModelType]):
         else:
             self.pid = os.getpid()      # process id
             print(f'[green]🏃‍♂️ {self}.on_startup() STARTUP (PROCESS)[/green]')
-        # abx.pm.hook.on_actor_startup(self)
+        # abx.pm.hook.on_actor_startup(actor=self)
         
-    def on_shutdown(self, err: BaseException | None=None) -> None:
+    def on_shutdown(self, last_obj: ModelType | None=None, last_error: BaseException | None=None) -> None:
-        print(f'[grey53]🏃‍♂️ {self}.on_shutdown() SHUTTING DOWN[/grey53]', err or '[green](gracefully)[/green]')
+        if isinstance(last_error, KeyboardInterrupt) or last_error is None:
-        # abx.pm.hook.on_actor_shutdown(self)
+            last_error_str = '[green](CTRL-C)[/green]'
+        elif isinstance(last_error, ActorQueueIsEmpty):
+            last_error_str = '[green](queue empty)[/green]'
+        elif isinstance(last_error, ActorObjectAlreadyClaimed):
+            last_error_str = '[green](queue race)[/green]'
+        else:
+            last_error_str = f'[red]{type(last_error).__name__}: {last_error}[/red]'
 
-    def on_tick_start(self, obj: ModelType) -> None:
+        print(f'[grey53]🏃‍♂️ {self}.on_shutdown() SHUTTING DOWN[/grey53] {last_error_str}')
-        # print(f'🏃‍♂️ {self}.on_tick_start()', obj.abid or obj.id)
+        # abx.pm.hook.on_actor_shutdown(actor=self, last_obj=last_obj, last_error=last_error)
-        # abx.pm.hook.on_actor_tick_start(self, obj_to_process)
+        
+    def on_tick_start(self, obj_to_process: ModelType) -> None:
+        print(f'🏃‍♂️ {self}.on_tick_start() {obj_to_process}')
+        # abx.pm.hook.on_actor_tick_start(actor=self, obj_to_process=obj)
         # self.timer = TimedProgress(self.MAX_TICK_TIME, prefix='      ')
-        pass
     
-    def on_tick_end(self, obj: ModelType) -> None:
+    def on_tick_end(self, obj_to_process: ModelType) -> None:
-        # print(f'🏃‍♂️ {self}.on_tick_end()', obj.abid or obj.id)
+        print(f'🏃‍♂️ {self}.on_tick_end() {obj_to_process}')
-        # abx.pm.hook.on_actor_tick_end(self, obj_to_process)
+        # abx.pm.hook.on_actor_tick_end(actor=self, obj_to_process=obj_to_process)
         # self.timer.end()
-        pass
     
-    def on_tick_exception(self, obj: ModelType, err: BaseException) -> None:
+    def on_tick_exception(self, obj_to_process: ModelType, error: Exception) -> None:
-        print(f'[red]🏃‍♂️ {self}.on_tick_exception()[/red]', obj.abid or obj.id, err)
+        print(f'[red]🏃‍♂️ {self}.on_tick_exception()[/red] {obj_to_process}: [red]{type(error).__name__}: {error}[/red]')
-        # abx.pm.hook.on_actor_tick_exception(self, obj_to_process, err)
+        # abx.pm.hook.on_actor_tick_exception(actor=self, obj_to_process=obj_to_process, error=error)
+
+
+
+def compile_sql_select(queryset: QuerySet, filter_kwargs: dict[str, Any] | None=None, order_args: tuple[str, ...]=(), limit: int | None=None) -> tuple[str, tuple[Any, ...]]:
+    """
+    Compute the SELECT query SQL for a queryset.filter(**filter_kwargs).order_by(*order_args)[:limit] call
+    Returns a tuple of (sql, params) where sql is a template string containing %s (unquoted) placeholders for the params
+    
+    WARNING:
+    final_sql = sql % params  DOES NOT WORK to assemble the final SQL string because the %s placeholders are not quoted/escaped
+    they should always passed separately to the DB driver so it can do its own quoting/escaping to avoid SQL injection and syntax errors
+    """
+    assert isinstance(queryset, QuerySet), f'compile_sql_select(...) first argument must be a QuerySet, got: {type(queryset).__name__} instead'
+    assert filter_kwargs is None or isinstance(filter_kwargs, dict), f'compile_sql_select(...) filter_kwargs argument must be a dict[str, Any], got: {type(filter_kwargs).__name__} instead'
+    assert isinstance(order_args, tuple) and all(isinstance(arg, str) for arg in order_args), f'compile_sql_select(...) order_args argument must be a tuple[str, ...] got: {type(order_args).__name__} instead'
+    assert limit is None or isinstance(limit, int), f'compile_sql_select(...) limit argument must be an int, got: {type(limit).__name__} instead'
+    
+    queryset = queryset._chain()                      # type: ignore   # copy queryset to avoid modifying the original
+    if filter_kwargs:
+        queryset = queryset.filter(**filter_kwargs)
+    if order_args:
+        queryset = queryset.order_by(*order_args)
+    if limit is not None:
+        queryset = queryset[:limit]
+    query = queryset.query
+    
+    # e.g. SELECT id FROM core_archiveresult WHERE status NOT IN (%s, %s, %s) AND retry_at <= %s ORDER BY retry_at ASC LIMIT 50
+    select_sql, select_params = query.get_compiler(queryset.db).as_sql()
+    return select_sql, select_params
+
+
+def compile_sql_update(queryset: QuerySet, update_kwargs: dict[str, Any], filter_kwargs: dict[str, Any] | None=None) -> tuple[str, tuple[Any, ...]]:
+    """
+    Compute the UPDATE query SQL for a queryset.filter(**filter_kwargs).update(**update_kwargs) call
+    Returns a tuple of (sql, params) where sql is a template string containing %s (unquoted) placeholders for the params
+    
+    Based on the django.db.models.QuerySet.update() source code, but modified to return the SQL instead of executing the update
+    https://github.com/django/django/blob/611bf6c2e2a1b4ab93273980c45150c099ab146d/django/db/models/query.py#L1217
+    
+    WARNING:
+    final_sql = sql % params  DOES NOT WORK to assemble the final SQL string because the %s placeholders are not quoted/escaped
+    they should always passed separately to the DB driver so it can do its own quoting/escaping to avoid SQL injection and syntax errors
+    """
+    assert isinstance(queryset, QuerySet), f'compile_sql_update(...) first argument must be a QuerySet, got: {type(queryset).__name__} instead'
+    assert isinstance(update_kwargs, dict), f'compile_sql_update(...) update_kwargs argument must be a dict[str, Any], got: {type(update_kwargs).__name__} instead'
+    assert filter_kwargs is None or isinstance(filter_kwargs, dict), f'compile_sql_update(...) filter_kwargs argument must be a dict[str, Any], got: {type(filter_kwargs).__name__} instead'
+    
+    queryset = queryset._chain()                      # type: ignore   # copy queryset to avoid modifying the original
+    if filter_kwargs:
+        queryset = queryset.filter(**filter_kwargs)
+    queryset.query.clear_ordering(force=True)                          # clear any ORDER BY clauses
+    queryset.query.clear_limits()                                      # clear any LIMIT clauses aka slices[:n]
+    queryset._for_write = True                        # type: ignore
+    query = queryset.query.chain(sql.UpdateQuery)     # type: ignore
+    query.add_update_values(update_kwargs)            # type: ignore
+    query.annotations = {}                                             # clear any annotations
+    
+    # e.g. UPDATE core_archiveresult SET status='%s', retry_at='%s' WHERE status NOT IN (%s, %s, %s) AND retry_at <= %s
+    update_sql, update_params = query.get_compiler(queryset.db).as_sql()
+    return update_sql, update_params

archivebox/actors/models.py

@@ -1,3 +1,300 @@
-from django.db import models
+from typing import ClassVar, Type, Iterable
+from datetime import datetime, timedelta
 
-# Create your models here.
+from statemachine.mixins import MachineMixin
+
+from django.db import models
+from django.utils import timezone
+from django.utils.functional import classproperty
+
+from statemachine import registry, StateMachine, State
+
+from django.core import checks
+
+class DefaultStatusChoices(models.TextChoices):
+    QUEUED = 'queued', 'Queued'
+    STARTED = 'started', 'Started'
+    SEALED = 'sealed', 'Sealed'
+
+
+default_status_field: models.CharField = models.CharField(choices=DefaultStatusChoices.choices, max_length=15, default=DefaultStatusChoices.QUEUED, null=False, blank=False, db_index=True)
+default_retry_at_field: models.DateTimeField = models.DateTimeField(default=timezone.now, null=False, db_index=True)
+
+ObjectState = State | str
+ObjectStateList = Iterable[ObjectState]
+
+
+class BaseModelWithStateMachine(models.Model, MachineMixin):
+    id: models.UUIDField
+    
+    StatusChoices: ClassVar[Type[models.TextChoices]]
+    
+    # status: models.CharField
+    # retry_at: models.DateTimeField
+    
+    state_machine_name: ClassVar[str]
+    state_field_name: ClassVar[str]
+    state_machine_attr: ClassVar[str] = 'sm'
+    bind_events_as_methods: ClassVar[bool] = True
+    
+    active_state: ClassVar[ObjectState]
+    retry_at_field_name: ClassVar[str]
+    
+    class Meta:
+        abstract = True
+        
+    @classmethod
+    def check(cls, sender=None, **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='actors.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='actors.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='actors.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 = {cls.id!r}',
+                obj=cls,
+                id='actors.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='actors.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='actors.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='actors.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='actors.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='actors.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)
+    
+    @RETRY_AT.setter
+    def RETRY_AT(self, value: datetime):
+        setattr(self, self.retry_at_field_name, value)
+        
+    @property
+    def STATE(self) -> str:
+        return getattr(self, self.state_field_name)
+    
+    @STATE.setter
+    def STATE(self, value: str):
+        setattr(self, self.state_field_name, value)
+        
+    def bump_retry_at(self, seconds: int = 10):
+        self.RETRY_AT = timezone.now() + timedelta(seconds=seconds)
+        
+    @classproperty
+    def ACTIVE_STATE(cls) -> str:
+        return cls._state_to_str(cls.StateMachineClass.active_state)
+        
+    @classproperty
+    def INITIAL_STATE(cls) -> str:
+        return cls._state_to_str(cls.StateMachineClass.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]
+        
+    @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__}'
+        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]
+            return models.TextChoices('StatusChoices', joined)
+        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
+
+    @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 that inherits from MachineMixin"""
+
+        model_state_machine_name = getattr(cls, 'state_machine_name', None)
+        if model_state_machine_name:
+            StateMachineCls = registry.get_machine_cls(model_state_machine_name)
+            assert issubclass(StateMachineCls, StateMachine)
+            return StateMachineCls
+        raise NotImplementedError(f'ActorType[{cls.__name__}] must define .state_machine_name: str that points to a valid StateMachine')
+    
+    # @classproperty
+    # def final_q(cls) -> Q:
+    #     """Get the filter for objects that are in a final state"""
+    #     return Q(**{f'{cls.state_field_name}__in': cls.final_states})
+    
+    # @classproperty
+    # def active_q(cls) -> Q:
+    #     """Get the filter for objects that are actively processing right now"""
+    #     return Q(**{cls.state_field_name: cls._state_to_str(cls.active_state)})   # e.g. Q(status='started')
+    
+    # @classproperty
+    # def stalled_q(cls) -> Q:
+    #     """Get the filter for objects that are marked active but have timed out"""
+    #     return cls.active_q & Q(retry_at__lte=timezone.now())                     # e.g. Q(status='started') AND Q(<retry_at is in the past>)
+    
+    # @classproperty
+    # def future_q(cls) -> Q:
+    #     """Get the filter for objects that have a retry_at in the future"""
+    #     return Q(retry_at__gt=timezone.now())
+    
+    # @classproperty
+    # def pending_q(cls) -> Q:
+    #     """Get the filter for objects that are ready for processing."""
+    #     return ~(cls.active_q) & ~(cls.final_q) & ~(cls.future_q)
+    
+    # @classmethod
+    # def get_queue(cls) -> QuerySet:
+    #     """
+    #     Get the sorted and filtered QuerySet of objects that are ready for processing.
+    #     e.g. qs.exclude(status__in=('sealed', 'started'), retry_at__gt=timezone.now()).order_by('retry_at')
+    #     """
+    #     return cls.objects.filter(cls.pending_q)
+
+
+class ModelWithStateMachine(BaseModelWithStateMachine):
+    StatusChoices: ClassVar[Type[DefaultStatusChoices]] = DefaultStatusChoices
+    
+    status: models.CharField = BaseModelWithStateMachine.StatusField()
+    retry_at: models.DateTimeField = BaseModelWithStateMachine.RetryAtField()
+    
+    state_machine_name: ClassVar[str]      # e.g. 'core.statemachines.ArchiveResultMachine'
+    state_field_name: ClassVar[str]        = 'status'
+    state_machine_attr: ClassVar[str]      = 'sm'
+    bind_events_as_methods: ClassVar[bool] = True
+    
+    active_state: ClassVar[str]            = StatusChoices.STARTED
+    retry_at_field_name: ClassVar[str]     = 'retry_at'
+    
+    class Meta:
+        abstract = True

archivebox/actors/orchestrator.py

@@ -6,7 +6,7 @@ import itertools
 from typing import Dict, Type, Literal, ClassVar
 from django.utils.functional import classproperty
 
-from multiprocessing import Process, cpu_count
+from multiprocessing import Process
 from threading import Thread, get_native_id
 
 

archivebox/actors/statemachine.py

@@ -1,286 +0,0 @@
-from statemachine import State, StateMachine
-from django.db import models
-from multiprocessing import Process
-import psutil
-import time
-
-# State Machine Definitions
-#################################################
-
-class SnapshotMachine(StateMachine):
-    """State machine for managing Snapshot lifecycle."""
-    
-    # States
-    queued = State(initial=True)
-    started = State()
-    sealed = State(final=True)
-    
-    # Transitions
-    start = queued.to(started, cond='can_start')
-    seal = started.to(sealed, cond='is_finished')
-    
-    # Events
-    tick = (
-        queued.to.itself(unless='can_start') |
-        queued.to(started, cond='can_start') |
-        started.to.itself(unless='is_finished') |
-        started.to(sealed, cond='is_finished')
-    )
-    
-    def __init__(self, snapshot):
-        self.snapshot = snapshot
-        super().__init__()
-        
-    def can_start(self):
-        return True
-        
-    def is_finished(self):
-        return not self.snapshot.has_pending_archiveresults()
-        
-    def before_start(self):
-        """Pre-start validation and setup."""
-        self.snapshot.cleanup_dir()
-        
-    def after_start(self):
-        """Post-start side effects."""
-        self.snapshot.create_pending_archiveresults()
-        self.snapshot.update_indices()
-        self.snapshot.bump_retry_at(seconds=10)
-        
-    def before_seal(self):
-        """Pre-seal validation and cleanup."""
-        self.snapshot.cleanup_dir()
-        
-    def after_seal(self):
-        """Post-seal actions."""
-        self.snapshot.update_indices()
-        self.snapshot.seal_dir()
-        self.snapshot.upload_dir()
-        self.snapshot.retry_at = None
-        self.snapshot.save()
-
-
-class ArchiveResultMachine(StateMachine):
-    """State machine for managing ArchiveResult lifecycle."""
-    
-    # States
-    queued = State(initial=True)
-    started = State()
-    succeeded = State(final=True)
-    backoff = State()
-    failed = State(final=True)
-    
-    # Transitions
-    start = queued.to(started, cond='can_start')
-    succeed = started.to(succeeded, cond='extractor_succeeded')
-    backoff = started.to(backoff, unless='extractor_succeeded')
-    retry = backoff.to(queued, cond='can_retry')
-    fail = backoff.to(failed, unless='can_retry')
-    
-    # Events
-    tick = (
-        queued.to.itself(unless='can_start') |
-        queued.to(started, cond='can_start') |
-        started.to.itself(cond='extractor_still_running') |
-        started.to(succeeded, cond='extractor_succeeded') |
-        started.to(backoff, unless='extractor_succeeded') |
-        backoff.to.itself(cond='still_waiting_to_retry') |
-        backoff.to(queued, cond='can_retry') |
-        backoff.to(failed, unless='can_retry')
-    )
-    
-    def __init__(self, archiveresult):
-        self.archiveresult = archiveresult
-        super().__init__()
-    
-    def can_start(self):
-        return True
-    
-    def extractor_still_running(self):
-        return self.archiveresult.start_ts > time.now() - timedelta(seconds=5)
-    
-    def extractor_succeeded(self):
-        # return check_if_extractor_succeeded(self.archiveresult)
-        return self.archiveresult.start_ts < time.now() - timedelta(seconds=5)
-    
-    def can_retry(self):
-        return self.archiveresult.retries < self.archiveresult.max_retries
-        
-    def before_start(self):
-        """Pre-start initialization."""
-        self.archiveresult.retries += 1
-        self.archiveresult.start_ts = time.now()
-        self.archiveresult.output = None
-        self.archiveresult.error = None
-        
-    def after_start(self):
-        """Post-start execution."""
-        self.archiveresult.bump_retry_at(seconds=self.archiveresult.timeout + 5)
-        execute_extractor(self.archiveresult)
-        self.archiveresult.snapshot.bump_retry_at(seconds=5)
-        
-    def before_succeed(self):
-        """Pre-success validation."""
-        self.archiveresult.output = get_archiveresult_output(self.archiveresult)
-        
-    def after_succeed(self):
-        """Post-success cleanup."""
-        self.archiveresult.end_ts = time.now()
-        self.archiveresult.retry_at = None
-        self.archiveresult.update_indices()
-        
-    def before_backoff(self):
-        """Pre-backoff error capture."""
-        self.archiveresult.error = get_archiveresult_error(self.archiveresult)
-        
-    def after_backoff(self):
-        """Post-backoff retry scheduling."""
-        self.archiveresult.end_ts = time.now()
-        self.archiveresult.bump_retry_at(
-            seconds=self.archiveresult.timeout * self.archiveresult.retries
-        )
-        self.archiveresult.update_indices()
-        
-    def before_fail(self):
-        """Pre-failure finalization."""
-        self.archiveresult.retry_at = None
-        
-    def after_fail(self):
-        """Post-failure cleanup."""
-        self.archiveresult.update_indices()
-
-# Models
-#################################################
-
-class Snapshot(models.Model):
-    status = models.CharField(max_length=32, default='queued')
-    retry_at = models.DateTimeField(null=True)
-    
-    @property
-    def sm(self):
-        """Get the state machine for this snapshot."""
-        return SnapshotMachine(self)
-    
-    def has_pending_archiveresults(self):
-        return self.archiveresult_set.exclude(
-            status__in=['succeeded', 'failed']
-        ).exists()
-    
-    def bump_retry_at(self, seconds):
-        self.retry_at = time.now() + timedelta(seconds=seconds)
-        self.save()
-        
-    def cleanup_dir(self):
-        cleanup_snapshot_dir(self)
-        
-    def create_pending_archiveresults(self):
-        create_snapshot_pending_archiveresults(self)
-        
-    def update_indices(self):
-        update_snapshot_index_json(self)
-        update_snapshot_index_html(self)
-        
-    def seal_dir(self):
-        seal_snapshot_dir(self)
-        
-    def upload_dir(self):
-        upload_snapshot_dir(self)
-
-
-class ArchiveResult(models.Model):
-    snapshot = models.ForeignKey(Snapshot, on_delete=models.CASCADE)
-    status = models.CharField(max_length=32, default='queued')
-    retry_at = models.DateTimeField(null=True)
-    retries = models.IntegerField(default=0)
-    max_retries = models.IntegerField(default=3)
-    timeout = models.IntegerField(default=60)
-    start_ts = models.DateTimeField(null=True)
-    end_ts = models.DateTimeField(null=True)
-    output = models.TextField(null=True)
-    error = models.TextField(null=True)
-    
-    def get_machine(self):
-        return ArchiveResultMachine(self)
-    
-    def bump_retry_at(self, seconds):
-        self.retry_at = time.now() + timedelta(seconds=seconds)
-        self.save()
-        
-    def update_indices(self):
-        update_archiveresult_index_json(self)
-        update_archiveresult_index_html(self)
-
-
-# Actor System
-#################################################
-
-class BaseActor:
-    MAX_TICK_TIME = 60
-    
-    def tick(self, obj):
-        """Process a single object through its state machine."""
-        machine = obj.get_machine()
-        
-        if machine.is_queued:
-            if machine.can_start():
-                machine.start()
-                
-        elif machine.is_started:
-            if machine.can_seal():
-                machine.seal()
-                
-        elif machine.is_backoff:
-            if machine.can_retry():
-                machine.retry()
-            else:
-                machine.fail()
-
-
-class Orchestrator:
-    """Main orchestrator that manages all actors."""
-    
-    def __init__(self):
-        self.pid = None
-        
-    @classmethod
-    def spawn(cls):
-        orchestrator = cls()
-        proc = Process(target=orchestrator.runloop)
-        proc.start()
-        return proc.pid
-        
-    def runloop(self):
-        self.pid = os.getpid()
-        abx.pm.hook.on_orchestrator_startup(self)
-        
-        try:
-            while True:
-                self.process_queue(Snapshot)
-                self.process_queue(ArchiveResult)
-                time.sleep(0.1)
-                
-        except (KeyboardInterrupt, SystemExit):
-            abx.pm.hook.on_orchestrator_shutdown(self)
-            
-    def process_queue(self, model):
-        retry_at_reached = Q(retry_at__isnull=True) | Q(retry_at__lte=time.now())
-        queue = model.objects.filter(retry_at_reached)
-        
-        if queue.exists():
-            actor = BaseActor()
-            for obj in queue:
-                try:
-                    with transaction.atomic():
-                        actor.tick(obj)
-                except Exception as e:
-                    abx.pm.hook.on_actor_tick_exception(actor, obj, e)
-
-
-# Periodic Tasks
-#################################################
-
-@djhuey.periodic_task(schedule=djhuey.crontab(minute='*'))
-def ensure_orchestrator_running():
-    """Ensure orchestrator is running, start if not."""
-    if not any(p.name().startswith('Orchestrator') for p in psutil.process_iter()):
-        Orchestrator.spawn()

archivebox/core/actors.py

@@ -2,72 +2,40 @@ __package__ = 'archivebox.core'
 
 from typing import ClassVar
 
-from rich import print
+from statemachine import State
-
-from django.db.models import QuerySet
-from django.utils import timezone
-from datetime import timedelta
-from core.models import Snapshot
 
+from core.models import Snapshot, ArchiveResult
+from core.statemachines import SnapshotMachine, ArchiveResultMachine
 from actors.actor import ActorType
 
 
 class SnapshotActor(ActorType[Snapshot]):
+    Model = Snapshot
+    StateMachineClass = SnapshotMachine
     
-    QUERYSET: ClassVar[QuerySet] = Snapshot.objects.filter(status='queued')
+    ACTIVE_STATE: ClassVar[State] = SnapshotMachine.started
-    CLAIM_WHERE: ClassVar[str] = 'status = "queued"'  # the WHERE clause to filter the objects when atomically getting the next object from the queue
+    FINAL_STATES: ClassVar[list[State]] = SnapshotMachine.final_states
-    CLAIM_SET: ClassVar[str] = 'status = "started"'   # the SET clause to claim the object when atomically getting the next object from the queue
+    STATE_FIELD_NAME: ClassVar[str] = SnapshotMachine.state_field_name
-    CLAIM_ORDER: ClassVar[str] = 'created_at DESC'    # the ORDER BY clause to sort the objects with when atomically getting the next object from the queue
-    CLAIM_FROM_TOP: ClassVar[int] = 50                # the number of objects to consider when atomically getting the next object from the queue
     
-    # model_type: Type[ModelType]
+    MAX_CONCURRENT_ACTORS: ClassVar[int] = 3
-    MAX_CONCURRENT_ACTORS: ClassVar[int] = 4               # min 2, max 8, up to 60% of available cpu cores
+    MAX_TICK_TIME: ClassVar[int] = 10
-    MAX_TICK_TIME: ClassVar[int] = 60                          # maximum duration in seconds to process a single object
+    CLAIM_FROM_TOP_N: ClassVar[int] = MAX_CONCURRENT_ACTORS * 10
 
-    def claim_sql_where(self) -> str:
-        """override this to implement a custom WHERE clause for the atomic claim step e.g. "status = 'queued' AND locked_by = NULL" """
-        return self.CLAIM_WHERE
 
-    def claim_sql_set(self) -> str:
-        """override this to implement a custom SET clause for the atomic claim step e.g. "status = 'started' AND locked_by = {self.pid}" """
-        retry_at = timezone.now() + timedelta(seconds=self.MAX_TICK_TIME)
-        # format as 2024-10-31 10:14:33.240903
-        retry_at_str = retry_at.strftime('%Y-%m-%d %H:%M:%S.%f')
-        return f'{self.CLAIM_SET}, retry_at = {retry_at_str}'
 
-    def claim_sql_order(self) -> str:
+class ArchiveResultActor(ActorType[ArchiveResult]):
-        """override this to implement a custom ORDER BY clause for the atomic claim step e.g. "created_at DESC" """
+    Model = ArchiveResult
-        return self.CLAIM_ORDER
+    StateMachineClass = ArchiveResultMachine
     
-    def claim_from_top(self) -> int:
+    ACTIVE_STATE: ClassVar[State] = ArchiveResultMachine.started
-        """override this to implement a custom number of objects to consider when atomically claiming the next object from the top of the queue"""
+    FINAL_STATES: ClassVar[list[State]] = ArchiveResultMachine.final_states
-        return self.CLAIM_FROM_TOP
+    STATE_FIELD_NAME: ClassVar[str] = ArchiveResultMachine.state_field_name
     
-    def tick(self, obj: Snapshot) -> None:
+    MAX_CONCURRENT_ACTORS: ClassVar[int] = 6
-        """override this to process the object"""
+    MAX_TICK_TIME: ClassVar[int] = 60
-        print(f'[blue]🏃‍♂️ {self}.tick()[/blue]', obj.abid or obj.id)
+    CLAIM_FROM_TOP_N: ClassVar[int] = MAX_CONCURRENT_ACTORS * 10
-        # For example:
-        # do_some_task(obj)
-        # do_something_else(obj)
-        # obj._model.objects.filter(pk=obj.pk, status='started').update(status='success')
-        # raise NotImplementedError('tick() must be implemented by the Actor subclass')
 
-    def on_shutdown(self, err: BaseException | None=None) -> None:
+    # @classproperty
-        print(f'[grey53]🏃‍♂️ {self}.on_shutdown() SHUTTING DOWN[/grey53]', err or '[green](gracefully)[/green]')
+    # def qs(cls) -> QuerySet[ModelType]:
-        # abx.pm.hook.on_actor_shutdown(self)
+    #     """Get the unfiltered and unsorted QuerySet of all objects that this Actor might care about."""
-        
+    #     return cls.Model.objects.filter(extractor='favicon')
-    def on_tick_start(self, obj: Snapshot) -> None:
-        # print(f'🏃‍♂️ {self}.on_tick_start()', obj.abid or obj.id)
-        # abx.pm.hook.on_actor_tick_start(self, obj_to_process)
-        # self.timer = TimedProgress(self.MAX_TICK_TIME, prefix='      ')
-        pass
-    
-    def on_tick_end(self, obj: Snapshot) -> None:
-        # print(f'🏃‍♂️ {self}.on_tick_end()', obj.abid or obj.id)
-        # abx.pm.hook.on_actor_tick_end(self, obj_to_process)
-        # self.timer.end()
-        pass
-    
-    def on_tick_exception(self, obj: Snapshot, err: BaseException) -> None:
-        print(f'[red]🏃‍♂️ {self}.on_tick_exception()[/red]', obj.abid or obj.id, err)
-        # abx.pm.hook.on_actor_tick_exception(self, obj_to_process, err)

archivebox/core/models.py

@@ -20,7 +20,7 @@ from django.db.models import Case, When, Value, IntegerField
 from django.contrib import admin
 from django.conf import settings
 
-from statemachine.mixins import MachineMixin
+from actors.models import ModelWithStateMachine
 
 from archivebox.config import CONSTANTS
 
@@ -156,7 +156,7 @@ class SnapshotManager(models.Manager):
         return super().get_queryset().prefetch_related('tags', 'archiveresult_set')  # .annotate(archiveresult_count=models.Count('archiveresult')).distinct()
 
 
-class Snapshot(ABIDModel, MachineMixin):
+class Snapshot(ABIDModel, ModelWithStateMachine):
     abid_prefix = 'snp_'
     abid_ts_src = 'self.created_at'
     abid_uri_src = 'self.url'
@@ -164,29 +164,27 @@ class Snapshot(ABIDModel, MachineMixin):
     abid_rand_src = 'self.id'
     abid_drift_allowed = True
 
-    state_field_name = 'status'
     state_machine_name = 'core.statemachines.SnapshotMachine'
-    state_machine_attr = 'sm'
+    state_field_name = 'status'
-    
+    retry_at_field_name = 'retry_at'
-    class SnapshotStatus(models.TextChoices):
+    StatusChoices = ModelWithStateMachine.StatusChoices
-        QUEUED = 'queued', 'Queued'
+    active_state = StatusChoices.STARTED
-        STARTED = 'started', 'Started'
-        SEALED = 'sealed', 'Sealed'
-        
-    status = models.CharField(max_length=15, default=SnapshotStatus.QUEUED, null=False, blank=False)
     
     id = models.UUIDField(primary_key=True, default=None, null=False, editable=False, unique=True, verbose_name='ID')
     abid = ABIDField(prefix=abid_prefix)
 
-    created_by = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, default=None, null=False, related_name='snapshot_set')
+    created_by = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, default=None, null=False, related_name='snapshot_set', db_index=True)
     created_at = AutoDateTimeField(default=None, null=False, db_index=True)  # loaded from self._init_timestamp
     modified_at = models.DateTimeField(auto_now=True)
     
+    status = ModelWithStateMachine.StatusField(choices=StatusChoices, default=StatusChoices.QUEUED)
+    retry_at = ModelWithStateMachine.RetryAtField(default=timezone.now)
+
     # legacy ts fields
     bookmarked_at = AutoDateTimeField(default=None, null=False, editable=True, db_index=True)
     downloaded_at = models.DateTimeField(default=None, null=True, editable=False, db_index=True, blank=True)
 
-    crawl = models.ForeignKey(Crawl, on_delete=models.CASCADE, default=None, null=True, blank=True, related_name='snapshot_set')
+    crawl: Crawl = models.ForeignKey(Crawl, on_delete=models.CASCADE, default=None, null=True, blank=True, related_name='snapshot_set', db_index=True)  # type: ignore
 
     url = models.URLField(unique=True, db_index=True)
     timestamp = models.CharField(max_length=32, unique=True, db_index=True, editable=False)
@@ -210,12 +208,14 @@ class Snapshot(ABIDModel, MachineMixin):
         return result
 
     def __repr__(self) -> str:
-        title = (self.title_stripped or '-')[:64]
+        url = self.url or '<no url set>'
-        return f'[{self.timestamp}] {self.url[:64]} ({title})'
+        created_at = self.created_at.strftime("%Y-%m-%d %H:%M") if self.created_at else '<no timestamp set>'
+        if self.id and self.url:
+            return f'[{self.ABID}] {url[:64]} @ {created_at}'
+        return f'[{self.abid_prefix}****not*saved*yet****] {url[:64]} @ {created_at}'
 
     def __str__(self) -> str:
-        title = (self.title_stripped or '-')[:64]
+        return repr(self)
-        return f'[{self.timestamp}] {self.url[:64]} ({title})'
 
     @classmethod
     def from_json(cls, info: dict):
@@ -413,8 +413,7 @@ class Snapshot(ABIDModel, MachineMixin):
         self.tags.add(*tags_id)
         
     def has_pending_archiveresults(self) -> bool:
-        pending_statuses = [ArchiveResult.ArchiveResultStatus.QUEUED, ArchiveResult.ArchiveResultStatus.STARTED]
+        pending_archiveresults = self.archiveresult_set.exclude(status__in=ArchiveResult.FINAL_OR_ACTIVE_STATES)
-        pending_archiveresults = self.archiveresult_set.filter(status__in=pending_statuses)
         return pending_archiveresults.exists()
     
     def create_pending_archiveresults(self) -> list['ArchiveResult']:
@@ -423,14 +422,11 @@ class Snapshot(ABIDModel, MachineMixin):
             archiveresult, _created = ArchiveResult.objects.get_or_create(
                 snapshot=self,
                 extractor=extractor,
-                status=ArchiveResult.ArchiveResultStatus.QUEUED,
+                status=ArchiveResult.INITIAL_STATE,
             )
             archiveresults.append(archiveresult)
         return archiveresults
 
-    def bump_retry_at(self, seconds: int = 10):
-        self.retry_at = timezone.now() + timedelta(seconds=seconds)
-
 
     # def get_storage_dir(self, create=True, symlink=True) -> Path:
     #     date_str = self.bookmarked_at.strftime('%Y%m%d')
@@ -479,7 +475,7 @@ class ArchiveResultManager(models.Manager):
             ).order_by('indexing_precedence')
         return qs
 
-class ArchiveResult(ABIDModel):
+class ArchiveResult(ABIDModel, ModelWithStateMachine):
     abid_prefix = 'res_'
     abid_ts_src = 'self.snapshot.created_at'
     abid_uri_src = 'self.snapshot.url'
@@ -487,19 +483,19 @@ class ArchiveResult(ABIDModel):
     abid_rand_src = 'self.id'
     abid_drift_allowed = True
     
-    state_field_name = 'status'
+    class StatusChoices(models.TextChoices):
+        QUEUED = 'queued', 'Queued'                     # pending, initial
+        STARTED = 'started', 'Started'                  # active
+        
+        BACKOFF = 'backoff', 'Waiting to retry'         # pending
+        SUCCEEDED = 'succeeded', 'Succeeded'            # final
+        FAILED = 'failed', 'Failed'                     # final
+        SKIPPED = 'skipped', 'Skipped'                  # final
+        
     state_machine_name = 'core.statemachines.ArchiveResultMachine'
-    state_machine_attr = 'sm'
+    retry_at_field_name = 'retry_at'
-
+    state_field_name = 'status'
-    class ArchiveResultStatus(models.TextChoices):
+    active_state = StatusChoices.STARTED
-        QUEUED = 'queued', 'Queued'
-        STARTED = 'started', 'Started'
-        SUCCEEDED = 'succeeded', 'Succeeded'
-        FAILED = 'failed', 'Failed'
-        SKIPPED = 'skipped', 'Skipped'
-        BACKOFF = 'backoff', 'Waiting to retry'
-        
-    status = models.CharField(max_length=15, choices=ArchiveResultStatus.choices, default=ArchiveResultStatus.QUEUED, null=False, blank=False)
 
     EXTRACTOR_CHOICES = (
         ('htmltotext', 'htmltotext'),
@@ -522,19 +518,22 @@ class ArchiveResult(ABIDModel):
     id = models.UUIDField(primary_key=True, default=None, null=False, editable=False, unique=True, verbose_name='ID')
     abid = ABIDField(prefix=abid_prefix)
 
-    created_by = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, default=None, null=False, related_name='archiveresult_set')
+    created_by = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, default=None, null=False, related_name='archiveresult_set', db_index=True)
     created_at = AutoDateTimeField(default=None, null=False, db_index=True)
     modified_at = models.DateTimeField(auto_now=True)
     
-    snapshot = models.ForeignKey(Snapshot, on_delete=models.CASCADE, to_field='id', db_column='snapshot_id')
+    status = ModelWithStateMachine.StatusField(choices=StatusChoices.choices, default=StatusChoices.QUEUED)
+    retry_at = ModelWithStateMachine.RetryAtField(default=timezone.now)
 
-    extractor = models.CharField(choices=EXTRACTOR_CHOICES, max_length=32)
+    snapshot: Snapshot = models.ForeignKey(Snapshot, on_delete=models.CASCADE)   # type: ignore
-    cmd = models.JSONField()
+
-    pwd = models.CharField(max_length=256)
+    extractor = models.CharField(choices=EXTRACTOR_CHOICES, max_length=32, blank=False, null=False, db_index=True)
+    cmd = models.JSONField(default=None, null=True, blank=True)
+    pwd = models.CharField(max_length=256, default=None, null=True, blank=True)
     cmd_version = models.CharField(max_length=128, default=None, null=True, blank=True)
-    output = models.CharField(max_length=1024)
+    output = models.CharField(max_length=1024, default=None, null=True, blank=True)
-    start_ts = models.DateTimeField(db_index=True)
+    start_ts = models.DateTimeField(default=None, null=True, blank=True)
-    end_ts = models.DateTimeField()
+    end_ts = models.DateTimeField(default=None, null=True, blank=True)
 
     # the network interface that was used to download this result
     # uplink = models.ForeignKey(NetworkInterface, on_delete=models.SET_NULL, null=True, blank=True, verbose_name='Network Interface Used')
@@ -545,10 +544,17 @@ class ArchiveResult(ABIDModel):
         verbose_name = 'Archive Result'
         verbose_name_plural = 'Archive Results Log'
 
+    def __repr__(self):
+        snapshot_id = getattr(self, 'snapshot_id', None)
+        url = self.snapshot.url if snapshot_id else '<no url set>'
+        created_at = self.snapshot.created_at.strftime("%Y-%m-%d %H:%M") if snapshot_id else '<no timestamp set>'
+        extractor = self.extractor or '<no extractor set>'
+        if self.id and snapshot_id:
+            return f'[{self.ABID}] {url[:64]} @ {created_at} -> {extractor}'
+        return f'[{self.abid_prefix}****not*saved*yet****] {url} @ {created_at} -> {extractor}'
 
     def __str__(self):
-        # return f'[{self.abid}] 📅 {self.start_ts.strftime("%Y-%m-%d %H:%M")} 📄 {self.extractor} {self.snapshot.url}'
+        return repr(self)
-        return self.extractor
 
     # TODO: finish connecting machine.models
     # @cached_property
@@ -559,6 +565,10 @@ class ArchiveResult(ABIDModel):
     def snapshot_dir(self):
         return Path(self.snapshot.link_dir)
     
+    @cached_property
+    def url(self):
+        return self.snapshot.url
+
     @property
     def api_url(self) -> str:
         # /api/v1/core/archiveresult/{uulid}
@@ -597,9 +607,6 @@ class ArchiveResult(ABIDModel):
     def output_exists(self) -> bool:
         return os.path.exists(self.output_path())
         
-    def bump_retry_at(self, seconds: int = 10):
-        self.retry_at = timezone.now() + timedelta(seconds=seconds)
-        
     def create_output_dir(self):
         snap_dir = self.snapshot_dir
         snap_dir.mkdir(parents=True, exist_ok=True)

archivebox/core/statemachines.py

@@ -16,9 +16,9 @@ class SnapshotMachine(StateMachine, strict_states=True):
     model: Snapshot
     
     # States
-    queued = State(value=Snapshot.SnapshotStatus.QUEUED, initial=True)
+    queued = State(value=Snapshot.StatusChoices.QUEUED, initial=True)
-    started = State(value=Snapshot.SnapshotStatus.STARTED)
+    started = State(value=Snapshot.StatusChoices.STARTED)
-    sealed = State(value=Snapshot.SnapshotStatus.SEALED, final=True)
+    sealed = State(value=Snapshot.StatusChoices.SEALED, final=True)
     
     # Tick Event
     tick = (
@@ -53,11 +53,11 @@ class ArchiveResultMachine(StateMachine, strict_states=True):
     model: ArchiveResult
     
     # States
-    queued = State(value=ArchiveResult.ArchiveResultStatus.QUEUED, initial=True)
+    queued = State(value=ArchiveResult.StatusChoices.QUEUED, initial=True)
-    started = State(value=ArchiveResult.ArchiveResultStatus.STARTED)
+    started = State(value=ArchiveResult.StatusChoices.STARTED)
-    backoff = State(value=ArchiveResult.ArchiveResultStatus.BACKOFF)
+    backoff = State(value=ArchiveResult.StatusChoices.BACKOFF)
-    succeeded = State(value=ArchiveResult.ArchiveResultStatus.SUCCEEDED, final=True)
+    succeeded = State(value=ArchiveResult.StatusChoices.SUCCEEDED, final=True)
-    failed = State(value=ArchiveResult.ArchiveResultStatus.FAILED, final=True)
+    failed = State(value=ArchiveResult.StatusChoices.FAILED, final=True)
     
     # Tick Event
     tick = (
@@ -78,7 +78,7 @@ class ArchiveResultMachine(StateMachine, strict_states=True):
         super().__init__(archiveresult, *args, **kwargs)
         
     def can_start(self) -> bool:
-        return self.archiveresult.snapshot and self.archiveresult.snapshot.is_started()
+        return self.archiveresult.snapshot and self.archiveresult.snapshot.STATE == Snapshot.active_state
     
     def is_succeeded(self) -> bool:
         return self.archiveresult.output_exists()
@@ -87,7 +87,10 @@ class ArchiveResultMachine(StateMachine, strict_states=True):
         return not self.archiveresult.output_exists()
     
     def is_backoff(self) -> bool:
-        return self.archiveresult.status == ArchiveResult.ArchiveResultStatus.BACKOFF
+        return self.archiveresult.STATE == ArchiveResult.StatusChoices.BACKOFF
+    
+    def is_finished(self) -> bool:
+        return self.is_failed() or self.is_succeeded()
 
     def on_started(self):
         self.archiveresult.start_ts = timezone.now()

archivebox/crawls/actors.py

@@ -2,68 +2,22 @@ __package__ = 'archivebox.crawls'
 
 from typing import ClassVar
 
-from rich import print
-
-from django.db.models import QuerySet
-
 from crawls.models import Crawl
+from crawls.statemachines import CrawlMachine
 
-from actors.actor import ActorType
+from actors.actor import ActorType, State
 
 
 class CrawlActor(ActorType[Crawl]):
+    """The Actor that manages the lifecycle of all Crawl objects"""
     
-    QUERYSET: ClassVar[QuerySet] = Crawl.objects.filter(status='queued')
+    Model = Crawl
-    CLAIM_WHERE: ClassVar[str] = 'status = "queued"'  # the WHERE clause to filter the objects when atomically getting the next object from the queue
+    StateMachineClass = CrawlMachine
-    CLAIM_SET: ClassVar[str] = 'status = "started"'   # the SET clause to claim the object when atomically getting the next object from the queue
-    CLAIM_ORDER: ClassVar[str] = 'created_at DESC'    # the ORDER BY clause to sort the objects with when atomically getting the next object from the queue
-    CLAIM_FROM_TOP: ClassVar[int] = 50                # the number of objects to consider when atomically getting the next object from the queue
     
-    # model_type: Type[ModelType]
+    ACTIVE_STATE: ClassVar[State] = CrawlMachine.started
-    MAX_CONCURRENT_ACTORS: ClassVar[int] = 4               # min 2, max 8, up to 60% of available cpu cores
+    FINAL_STATES: ClassVar[list[State]] = CrawlMachine.final_states
-    MAX_TICK_TIME: ClassVar[int] = 60                          # maximum duration in seconds to process a single object
+    STATE_FIELD_NAME: ClassVar[str] = Crawl.state_field_name
     
-    def claim_sql_where(self) -> str:
+    MAX_CONCURRENT_ACTORS: ClassVar[int] = 3
-        """override this to implement a custom WHERE clause for the atomic claim step e.g. "status = 'queued' AND locked_by = NULL" """
+    MAX_TICK_TIME: ClassVar[int] = 10
-        return self.CLAIM_WHERE
+    CLAIM_FROM_TOP_N: ClassVar[int] = MAX_CONCURRENT_ACTORS * 10
-    
-    def claim_sql_set(self) -> str:
-        """override this to implement a custom SET clause for the atomic claim step e.g. "status = 'started' AND locked_by = {self.pid}" """
-        return self.CLAIM_SET
-    
-    def claim_sql_order(self) -> str:
-        """override this to implement a custom ORDER BY clause for the atomic claim step e.g. "created_at DESC" """
-        return self.CLAIM_ORDER
-    
-    def claim_from_top(self) -> int:
-        """override this to implement a custom number of objects to consider when atomically claiming the next object from the top of the queue"""
-        return self.CLAIM_FROM_TOP
-        
-    def tick(self, obj: Crawl) -> None:
-        """override this to process the object"""
-        print(f'[blue]🏃‍♂️ {self}.tick()[/blue]', obj.abid or obj.id)
-        # For example:
-        # do_some_task(obj)
-        # do_something_else(obj)
-        # obj._model.objects.filter(pk=obj.pk, status='started').update(status='success')
-        # raise NotImplementedError('tick() must be implemented by the Actor subclass')
-    
-    def on_shutdown(self, err: BaseException | None=None) -> None:
-        print(f'[grey53]🏃‍♂️ {self}.on_shutdown() SHUTTING DOWN[/grey53]', err or '[green](gracefully)[/green]')
-        # abx.pm.hook.on_actor_shutdown(self)
-        
-    def on_tick_start(self, obj: Crawl) -> None:
-        # print(f'🏃‍♂️ {self}.on_tick_start()', obj.abid or obj.id)
-        # abx.pm.hook.on_actor_tick_start(self, obj_to_process)
-        # self.timer = TimedProgress(self.MAX_TICK_TIME, prefix='      ')
-        pass
-    
-    def on_tick_end(self, obj: Crawl) -> None:
-        # print(f'🏃‍♂️ {self}.on_tick_end()', obj.abid or obj.id)
-        # abx.pm.hook.on_actor_tick_end(self, obj_to_process)
-        # self.timer.end()
-        pass
-    
-    def on_tick_exception(self, obj: Crawl, err: BaseException) -> None:
-        print(f'[red]🏃‍♂️ {self}.on_tick_exception()[/red]', obj.abid or obj.id, err)
-        # abx.pm.hook.on_actor_tick_exception(self, obj_to_process, err)

archivebox/crawls/models.py

@@ -11,7 +11,7 @@ from django.conf import settings
 from django.urls import reverse_lazy
 from django.utils import timezone
 
-from statemachine.mixins import MachineMixin
+from actors.models import ModelWithStateMachine
 
 if TYPE_CHECKING:
     from core.models import Snapshot
@@ -50,7 +50,7 @@ class CrawlSchedule(ABIDModel, ModelWithHealthStats):
 
     
 
-class Crawl(ABIDModel, ModelWithHealthStats, MachineMixin):
+class Crawl(ABIDModel, ModelWithHealthStats, ModelWithStateMachine):
     """
     A single session of URLs to archive starting from a given Seed and expanding outwards. An "archiving session" so to speak.
 
@@ -67,17 +67,11 @@ class Crawl(ABIDModel, ModelWithHealthStats, MachineMixin):
     abid_rand_src = 'self.id'
     abid_drift_allowed = True
     
-    state_field_name = 'status'
     state_machine_name = 'crawls.statemachines.CrawlMachine'
-    state_machine_attr = 'sm'
+    retry_at_field_name = 'retry_at'
-    bind_events_as_methods = True
+    state_field_name = 'status'
-
+    StatusChoices = ModelWithStateMachine.StatusChoices
-    class CrawlStatus(models.TextChoices):
+    active_state = StatusChoices.STARTED
-        QUEUED = 'queued', 'Queued'
-        STARTED = 'started', 'Started'
-        SEALED = 'sealed', 'Sealed'
-
-    status = models.CharField(choices=CrawlStatus.choices, max_length=15, default=CrawlStatus.QUEUED, null=False, blank=False)
     
     id = models.UUIDField(primary_key=True, default=None, null=False, editable=False, unique=True, verbose_name='ID')
     abid = ABIDField(prefix=abid_prefix)
@@ -86,6 +80,8 @@ class Crawl(ABIDModel, ModelWithHealthStats, MachineMixin):
     created_at = AutoDateTimeField(default=None, null=False, db_index=True)
     modified_at = models.DateTimeField(auto_now=True)
     
+    status = ModelWithStateMachine.StatusField(choices=StatusChoices, default=StatusChoices.QUEUED)
+    retry_at = ModelWithStateMachine.RetryAtField(default=timezone.now)
 
     seed = models.ForeignKey(Seed, on_delete=models.PROTECT, related_name='crawl_set', null=False, blank=False)
     max_depth = models.PositiveSmallIntegerField(default=0, validators=[MinValueValidator(0), MaxValueValidator(4)])
@@ -127,10 +123,8 @@ class Crawl(ABIDModel, ModelWithHealthStats, MachineMixin):
     def has_pending_archiveresults(self) -> bool:
         from core.models import ArchiveResult
         
-        pending_statuses = [ArchiveResult.ArchiveResultStatus.QUEUED, ArchiveResult.ArchiveResultStatus.STARTED]
-        
         snapshot_ids = self.snapshot_set.values_list('id', flat=True)
-        pending_archiveresults = ArchiveResult.objects.filter(snapshot_id__in=snapshot_ids, status__in=pending_statuses)
+        pending_archiveresults = ArchiveResult.objects.filter(snapshot_id__in=snapshot_ids).exclude(status__in=ArchiveResult.FINAL_OR_ACTIVE_STATES)
         return pending_archiveresults.exists()
     
     def create_root_snapshot(self) -> 'Snapshot':
@@ -139,13 +133,10 @@ class Crawl(ABIDModel, ModelWithHealthStats, MachineMixin):
         root_snapshot, _ = Snapshot.objects.get_or_create(
             crawl=self,
             url=self.seed.uri,
+            status=Snapshot.INITIAL_STATE,
         )
         return root_snapshot
 
-    def bump_retry_at(self, seconds: int = 10):
-        self.retry_at = timezone.now() + timedelta(seconds=seconds)
-        self.save()
-
 
 class Outlink(models.Model):
     """A record of a link found on a page, pointing to another page."""

archivebox/crawls/statemachines.py

@@ -14,9 +14,9 @@ class CrawlMachine(StateMachine, strict_states=True):
     model: Crawl
     
     # States
-    queued = State(value=Crawl.CrawlStatus.QUEUED, initial=True)
+    queued = State(value=Crawl.StatusChoices.QUEUED, initial=True)
-    started = State(value=Crawl.CrawlStatus.STARTED)
+    started = State(value=Crawl.StatusChoices.STARTED)
-    sealed = State(value=Crawl.CrawlStatus.SEALED, final=True)
+    sealed = State(value=Crawl.StatusChoices.SEALED, final=True)
     
     # Tick Event
     tick = (