class Trigger(Iterator[datetime], metaclass=ABCMeta):
class Serializer(metaclass=ABCMeta):
class Subscription(metaclass=ABCMeta):
class EventSource(metaclass=ABCMeta):
class EventBroker(EventSource):
class AsyncEventBroker(EventSource):
class DataStore:
class AsyncDataStore:
    def next(self) -> datetime | None:
    def __getstate__(self):
    def __setstate__(self, state):
    def __iter__(self):
    def __next__(self) -> datetime:
    def serialize(self, obj) -> bytes:
    def serialize_to_unicode(self, obj) -> str:
    def deserialize(self, serialized: bytes):
    def deserialize_from_unicode(self, serialized: str):
    def __enter__(self) -> Subscription:
    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
    def unsubscribe(self) -> None:
    def subscribe(
        self, callback: Callable[[Event], Any],
        event_types: Iterable[type[Event]] | None = None,
        *,
        one_shot: bool = False
    ) -> Subscription:
    def __enter__(self):
    def __exit__(self, exc_type, exc_val, exc_tb):
    def publish(self, event: Event) -> None:
    def publish_local(self, event: Event) -> None:
    async def __aenter__(self):
    async def __aexit__(self, exc_type, exc_val, exc_tb):
    async def publish(self, event: Event) -> None:
    async def publish_local(self, event: Event) -> None:
    def __enter__(self):
    def __exit__(self, exc_type, exc_val, exc_tb):
    def events(self) -> EventSource:
    def add_task(self, task: Task) -> None:
    def remove_task(self, task_id: str) -> None:
    def get_task(self, task_id: str) -> Task:
    def get_tasks(self) -> list[Task]:
    def get_schedules(self, ids: set[str] | None = None) -> list[Schedule]:
    def add_schedule(self, schedule: Schedule, conflict_policy: ConflictPolicy) -> None:
    def remove_schedules(self, ids: Iterable[str]) -> None:
    def acquire_schedules(self, scheduler_id: str, limit: int) -> list[Schedule]:
    def release_schedules(self, scheduler_id: str, schedules: list[Schedule]) -> None:
    def get_next_schedule_run_time(self) -> datetime | None:
    def add_job(self, job: Job) -> None:
    def get_jobs(self, ids: Iterable[UUID] | None = None) -> list[Job]:
    def acquire_jobs(self, worker_id: str, limit: int | None = None) -> list[Job]:
    def release_job(self, worker_id: str, task_id: str, result: JobResult) -> None:
    def get_job_result(self, job_id: UUID) -> JobResult | None:
    async def __aenter__(self):
    async def __aexit__(self, exc_type, exc_val, exc_tb):
    def events(self) -> EventSource:
    async def add_task(self, task: Task) -> None:
    async def remove_task(self, task_id: str) -> None:
    async def get_task(self, task_id: str) -> Task:
    async def get_tasks(self) -> list[Task]:
    async def get_schedules(self, ids: set[str] | None = None) -> list[Schedule]:
    async def add_schedule(self, schedule: Schedule, conflict_policy: ConflictPolicy) -> None:
    async def remove_schedules(self, ids: Iterable[str]) -> None:
    async def acquire_schedules(self, scheduler_id: str, limit: int) -> list[Schedule]:
    async def release_schedules(self, scheduler_id: str, schedules: list[Schedule]) -> None:
    async def get_next_schedule_run_time(self) -> datetime | None:
    async def add_job(self, job: Job) -> None:
    async def get_jobs(self, ids: Iterable[UUID] | None = None) -> list[Job]:
    async def acquire_jobs(self, worker_id: str, limit: int | None = None) -> list[Job]:
    async def release_job(self, worker_id: str, task_id: str, result: JobResult) -> None:
    async def get_job_result(self, job_id: UUID) -> JobResult | None: def as_async_datastore(value: abc.DataStore | abc.AsyncDataStore) -> abc.AsyncDataStore:
    if isinstance(value, abc.DataStore):
        from apscheduler.datastores.async_adapter import AsyncDataStoreAdapter
        return AsyncDataStoreAdapter(value)

    return value class JobOutcome(Enum):
    success = auto()
    error = auto()
    missed_start_deadline = auto()
    cancelled = auto()
    expired = auto() class RunState(Enum):
    starting = auto()
    started = auto()
    stopping = auto()
    stopped = auto() class LocalAsyncEventBroker(AsyncEventBroker, BaseEventBroker):
    _task_group: TaskGroup = attrs.field(init=False)
    _exit_stack: AsyncExitStack = attrs.field(init=False)

    async def __aenter__(self) -> LocalAsyncEventBroker:
        self._exit_stack = AsyncExitStack()

        self._task_group = create_task_group()
        await self._exit_stack.enter_async_context(self._task_group)

        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
        del self._task_group

    async def publish(self, event: Event) -> None:
        await self.publish_local(event)

    async def publish_local(self, event: Event) -> None:
        event_type = type(event)
        one_shot_tokens: list[object] = []
        for _token, subscription in self._subscriptions.items():
            if subscription.event_types is None or event_type in subscription.event_types:
                self._task_group.start_soon(self._deliver_event, subscription.callback, event)
                if subscription.one_shot:
                    one_shot_tokens.append(subscription.token)

        for token in one_shot_tokens:
            super().unsubscribe(token)

    async def _deliver_event(self, func: Callable[[Event], Any], event: Event) -> None:
        try:
            retval = func(event)
            if iscoroutine(retval):
                await retval
        except BaseException:
            self._logger.exception('Error delivering %s event', event.__class__.__name__) class JobAdded(DataStoreEvent):
    job_id: UUID = attrs.field(converter=as_uuid)
    task_id: str
    schedule_id: str | None
    tags: frozenset[str] = attrs.field(converter=frozenset) class WorkerStarted(WorkerEvent):
    pass class WorkerStopped(WorkerEvent):
    exception: BaseException | None = None class JobInfo:
    job_id: UUID
    task_id: str
    schedule_id: str | None
    scheduled_fire_time: datetime | None
    jitter: timedelta
    start_deadline: datetime | None
    tags: frozenset[str]

    @classmethod
    def from_job(cls, job: Job) -> JobInfo:
        return cls(job_id=job.id, task_id=job.task_id, schedule_id=job.schedule_id,
                   scheduled_fire_time=job.scheduled_fire_time, jitter=job.jitter,
                   start_deadline=job.start_deadline, tags=job.tags) class JobResult:
    job_id: UUID
    outcome: JobOutcome = attrs.field(eq=False, order=False, converter=as_enum(JobOutcome))
    finished_at: datetime = attrs.field(eq=False, order=False,
                                        factory=partial(datetime.now, timezone.utc))
    exception: BaseException | None = attrs.field(eq=False, order=False, default=None)
    return_value: Any = attrs.field(eq=False, order=False, default=None)

    def marshal(self, serializer: abc.Serializer) -> dict[str, Any]:
        marshalled = attrs.asdict(self, value_serializer=serialize)
        if self.outcome is JobOutcome.error:
            marshalled['exception'] = serializer.serialize(self.exception)
        else:
            del marshalled['exception']

        if self.outcome is JobOutcome.success:
            marshalled['return_value'] = serializer.serialize(self.return_value)
        else:
            del marshalled['return_value']

        return marshalled

    @classmethod
    def unmarshal(cls, serializer: abc.Serializer, marshalled: dict[str, Any]) -> JobResult:
        if marshalled.get('exception'):
            marshalled['exception'] = serializer.deserialize(marshalled['exception'])
        elif marshalled.get('return_value'):
            marshalled['return_value'] = serializer.deserialize(marshalled['return_value'])

        return cls(**marshalled) def positive_integer(inst, field: attrs.Attribute, value) -> None:
    if value <= 0:
        raise ValueError(f'{field} must be a positive integer')
import os
import platform
import anyio
import attrs
from contextlib import AsyncExitStack
from datetime import datetime, timezone
from inspect import isawaitable
from logging import Logger, getLogger
from typing import Callable
from uuid import UUID
from anyio import TASK_STATUS_IGNORED, create_task_group, get_cancelled_exc_class, move_on_after
from anyio.abc import CancelScope
from ..abc import AsyncDataStore, EventSource, Job
from ..context import current_worker, job_info
from ..converters import as_async_datastore
from ..enums import JobOutcome, RunState
from ..eventbrokers.async_local import LocalAsyncEventBroker
from ..events import JobAdded, WorkerStarted, WorkerStopped
from ..structures import JobInfo, JobResult
from ..validators import positive_integer
from __future__ import annotations





@attrs.define(eq=False)
class AsyncWorker:
    """Runs jobs locally in a task group."""
    data_store: AsyncDataStore = attrs.field(converter=as_async_datastore)
    max_concurrent_jobs: int = attrs.field(kw_only=True, validator=positive_integer, default=100)
    identity: str = attrs.field(kw_only=True, default=None)
    logger: Logger | None = attrs.field(kw_only=True, default=getLogger(__name__))

    _state: RunState = attrs.field(init=False, default=RunState.stopped)
    _wakeup_event: anyio.Event = attrs.field(init=False, factory=anyio.Event)
    _acquired_jobs: set[Job] = attrs.field(init=False, factory=set)
    _events: LocalAsyncEventBroker = attrs.field(init=False, factory=LocalAsyncEventBroker)
    _running_jobs: set[UUID] = attrs.field(init=False, factory=set)
    _exit_stack: AsyncExitStack = attrs.field(init=False)

    def __attrs_post_init__(self) -> None:
        if not self.identity:
            self.identity = f'{platform.node()}-{os.getpid()}-{id(self)}'

    @property
    def events(self) -> EventSource:
        return self._events

    @property
    def state(self) -> RunState:
        return self._state

    async def __aenter__(self):
        self._state = RunState.starting
        self._wakeup_event = anyio.Event()
        self._exit_stack = AsyncExitStack()
        await self._exit_stack.__aenter__()
        await self._exit_stack.enter_async_context(self._events)

        # Initialize the data store and start relaying events to the worker's event broker
        await self._exit_stack.enter_async_context(self.data_store)
        self._exit_stack.enter_context(self.data_store.events.subscribe(self._events.publish))

        # Wake up the worker if the data store emits a significant job event
        self._exit_stack.enter_context(
            self.data_store.events.subscribe(lambda event: self._wakeup_event.set(), {JobAdded})
        )

        # Start the actual worker
        task_group = create_task_group()
        await self._exit_stack.enter_async_context(task_group)
        await task_group.start(self.run)
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        self._state = RunState.stopping
        self._wakeup_event.set()
        await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
        del self._wakeup_event

    async def run(self, *, task_status=TASK_STATUS_IGNORED) -> None:
        if self._state is not RunState.starting:
            raise RuntimeError(f'This function cannot be called while the worker is in the '
                               f'{self._state} state')

        # Set the current worker
        token = current_worker.set(self)

        # Signal that the worker has started
        self._state = RunState.started
        task_status.started()
        await self._events.publish(WorkerStarted())

        exception: BaseException | None = None
        try:
            async with create_task_group() as tg:
                while self._state is RunState.started:
                    limit = self.max_concurrent_jobs - len(self._running_jobs)
                    jobs = await self.data_store.acquire_jobs(self.identity, limit)
                    for job in jobs:
                        task = await self.data_store.get_task(job.task_id)
                        self._running_jobs.add(job.id)
                        tg.start_soon(self._run_job, job, task.func)

                    await self._wakeup_event.wait()
                    self._wakeup_event = anyio.Event()
        except get_cancelled_exc_class():
            pass
        except BaseException as exc:
            self.logger.exception('Worker crashed')
            exception = exc
        else:
            self.logger.info('Worker stopped')
        finally:

            current_worker.reset(token)