class DataStore:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    @property
    @abstractmethod
    def events(self) -> EventSource:
        pass

    @abstractmethod
    def add_task(self, task: Task) -> None:
        """
        Add the given task to the store.

        If a task with the same ID already exists, it replaces the old one but does NOT affect
        task accounting (# of running jobs).

        :param task: the task to be added
        """

    @abstractmethod
    def remove_task(self, task_id: str) -> None:
        """
        Remove the task with the given ID.

        :param task_id: ID of the task to be removed
        :raises TaskLookupError: if no matching task was found
        """

    @abstractmethod
    def get_task(self, task_id: str) -> Task:
        """
        Get an existing task definition.

        :param task_id: ID of the task to be returned
        :return: the matching task
        :raises TaskLookupError: if no matching task was found
        """

    @abstractmethod
    def get_tasks(self) -> list[Task]:
        """
        Get all the tasks in this store.

        :return: a list of tasks, sorted by ID
        """

    @abstractmethod
    def get_schedules(self, ids: set[str] | None = None) -> list[Schedule]:
        """
        Get schedules from the data store.

        :param ids: a specific set of schedule IDs to return, or ``None`` to return all schedules
        :return: the list of matching schedules, in unspecified order
        """

    @abstractmethod
    def add_schedule(self, schedule: Schedule, conflict_policy: ConflictPolicy) -> None:
        """
        Add or update the given schedule in the data store.

        :param schedule: schedule to be added
        :param conflict_policy: policy that determines what to do if there is an existing schedule
            with the same ID
        """

    @abstractmethod
    def remove_schedules(self, ids: Iterable[str]) -> None:
        """
        Remove schedules from the data store.

        :param ids: a specific set of schedule IDs to remove
        """

    @abstractmethod
    def acquire_schedules(self, scheduler_id: str, limit: int) -> list[Schedule]:
        """
        Acquire unclaimed due schedules for processing.

        This method claims up to the requested number of schedules for the given scheduler and
        returns them.

        :param scheduler_id: unique identifier of the scheduler
        :param limit: maximum number of schedules to claim
        :return: the list of claimed schedules
        """

    @abstractmethod
    def release_schedules(self, scheduler_id: str, schedules: list[Schedule]) -> None:
        """
        Release the claims on the given schedules and update them on the store.

        :param scheduler_id: unique identifier of the scheduler
        :param schedules: the previously claimed schedules
        """

    @abstractmethod
    def get_next_schedule_run_time(self) -> datetime | None:
        """
        Return the earliest upcoming run time of all the schedules in the store, or ``None`` if
        there are no active schedules.
        """

    @abstractmethod
    def add_job(self, job: Job) -> None:
        """
        Add a job to be executed by an eligible worker.

        :param job: the job object
        """

    @abstractmethod
    def get_jobs(self, ids: Iterable[UUID] | None = None) -> list[Job]:
        """
        Get the list of pending jobs.

        :param ids: a specific set of job IDs to return, or ``None`` to return all jobs
        :return: the list of matching pending jobs, in the order they will be given to workers
        """

    @abstractmethod
    def acquire_jobs(self, worker_id: str, limit: int | None = None) -> list[Job]:
        """
        Acquire unclaimed jobs for execution.

        This method claims up to the requested number of jobs for the given worker and returns
        them.

        :param worker_id: unique identifier of the worker
        :param limit: maximum number of jobs to claim and return
        :return: the list of claimed jobs
        """

    @abstractmethod
    def release_job(self, worker_id: str, task_id: str, result: JobResult) -> None:
        """
        Release the claim on the given job and record the result.

        :param worker_id: unique identifier of the worker
        :param task_id: the job's task ID
        :param result: the result of the job
        """

    @abstractmethod
    def get_job_result(self, job_id: UUID) -> JobResult | None:
        """
        Retrieve the result of a job.

        The result is removed from the store after retrieval.

        :param job_id: the identifier of the job
        :return: the result, or ``None`` if the result was not found
        """ class EventSource(metaclass=ABCMeta):
    """Interface for objects that can deliver notifications to interested subscribers."""

    @abstractmethod
    def subscribe(
        self, callback: Callable[[Event], Any],
        event_types: Iterable[type[Event]] | None = None,
        *,
        one_shot: bool = False
    ) -> Subscription:
        """
        Subscribe to events from this event source.

        :param callback: callable to be called with the event object when an event is published
        :param event_types: an iterable of concrete Event classes to subscribe to
        :param one_shot: if ``True``, automatically unsubscribe after the first matching event
        """ 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 LocalEventBroker(BaseEventBroker):
    _executor: ThreadPoolExecutor = attrs.field(init=False)
    _exit_stack: ExitStack = attrs.field(init=False)
    _subscriptions_lock: Lock = attrs.field(init=False, factory=Lock)

    def __enter__(self):
        self._exit_stack = ExitStack()
        self._executor = self._exit_stack.enter_context(ThreadPoolExecutor(1))
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._exit_stack.__exit__(exc_type, exc_val, exc_tb)
        del self._executor

    def subscribe(self, callback: Callable[[Event], Any],
                  event_types: Iterable[type[Event]] | None = None, *,
                  one_shot: bool = False) -> Subscription:
        if iscoroutinefunction(callback):
            raise ValueError('Coroutine functions are not supported as callbacks on a synchronous '
                             'event source')

        with self._subscriptions_lock:
            return super().subscribe(callback, event_types, one_shot=one_shot)

    def unsubscribe(self, token: object) -> None:
        with self._subscriptions_lock:
            super().unsubscribe(token)

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

    def publish_local(self, event: Event) -> None:
        event_type = type(event)
        with self._subscriptions_lock:
            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._executor.submit(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)

    def _deliver_event(self, func: Callable[[Event], Any], event: Event) -> None:
        try:
            func(event)
        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 Job:
    id: UUID = attrs.field(factory=uuid4)
    task_id: str = attrs.field(eq=False, order=False)
    args: tuple = attrs.field(eq=False, order=False, converter=tuple, default=())
    kwargs: dict[str, Any] = attrs.field(eq=False, order=False, converter=dict, default=())
    schedule_id: str | None = attrs.field(eq=False, order=False, default=None)
    scheduled_fire_time: datetime | None = attrs.field(eq=False, order=False, default=None)
    jitter: timedelta = attrs.field(eq=False, order=False, converter=as_timedelta,
                                    factory=timedelta)
    start_deadline: datetime | None = attrs.field(eq=False, order=False, default=None)
    tags: frozenset[str] = attrs.field(eq=False, order=False, converter=frozenset, default=())
    created_at: datetime = attrs.field(eq=False, order=False,
                                       factory=partial(datetime.now, timezone.utc))
    started_at: datetime | None = attrs.field(eq=False, order=False, default=None)
    acquired_by: str | None = attrs.field(eq=False, order=False, default=None)
    acquired_until: datetime | None = attrs.field(eq=False, order=False, default=None)

    @property
    def original_scheduled_time(self) -> datetime | None:
        """The scheduled time without any jitter included."""
        if self.scheduled_fire_time is None:
            return None

        return self.scheduled_fire_time - self.jitter

    def marshal(self, serializer: abc.Serializer) -> dict[str, Any]:
        marshalled = attrs.asdict(self, value_serializer=serialize)
        marshalled['args'] = serializer.serialize(self.args)
        marshalled['kwargs'] = serializer.serialize(self.kwargs)
        if not self.acquired_by:
            del marshalled['acquired_by']
            del marshalled['acquired_until']

        return marshalled

    @classmethod
    def unmarshal(cls, serializer: abc.Serializer, marshalled: dict[str, Any]) -> Job:
        marshalled['args'] = serializer.deserialize(marshalled['args'])
        marshalled['kwargs'] = serializer.deserialize(marshalled['kwargs'])
        return cls(**marshalled) 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 threading
import attrs
from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
from contextlib import ExitStack
from contextvars import copy_context
from datetime import datetime, timezone
from logging import Logger, getLogger
from typing import Callable
from uuid import UUID
from ..abc import DataStore, EventSource
from ..context import current_worker, job_info
from ..enums import JobOutcome, RunState
from ..eventbrokers.local import LocalEventBroker
from ..events import JobAdded, WorkerStarted, WorkerStopped
from ..structures import Job, JobInfo, JobResult
from ..validators import positive_integer
from __future__ import annotations





@attrs.define(eq=False)
class Worker:
    """Runs jobs locally in a thread pool."""
    data_store: DataStore
    max_concurrent_jobs: int = attrs.field(kw_only=True, validator=positive_integer, default=20)
    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: threading.Event = attrs.field(init=False)
    _acquired_jobs: set[Job] = attrs.field(init=False, factory=set)
    _events: LocalEventBroker = attrs.field(init=False, factory=LocalEventBroker)
    _running_jobs: set[UUID] = attrs.field(init=False, factory=set)
    _exit_stack: ExitStack = attrs.field(init=False)
    _executor: ThreadPoolExecutor = 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

    def __enter__(self) -> Worker:
        self._state = RunState.starting
        self._wakeup_event = threading.Event()
        self._exit_stack = ExitStack()
        self._exit_stack.__enter__()
        self._exit_stack.enter_context(self._events)

        # Initialize the data store and start relaying events to the worker's event broker
        self._exit_stack.enter_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 worker and return when it has signalled readiness or raised an exception
        start_future: Future[None] = Future()
        with self._events.subscribe(start_future.set_result, one_shot=True):
            self._executor = ThreadPoolExecutor(1)
            run_future = self._executor.submit(copy_context().run, self.run)
            wait([start_future, run_future], return_when=FIRST_COMPLETED)

        if run_future.done():
            run_future.result()

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._state = RunState.stopping
        self._wakeup_event.set()
        self._executor.shutdown(wait=exc_type is None)
        self._exit_stack.__exit__(exc_type, exc_val, exc_tb)
        del self._wakeup_event

    def run(self) -> 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
        self._events.publish(WorkerStarted())

        executor = ThreadPoolExecutor(max_workers=self.max_concurrent_jobs)
        exception: BaseException | None = None
        try:
            while self._state is RunState.started:
                available_slots = self.max_concurrent_jobs - len(self._running_jobs)
                if available_slots:
                    jobs = self.data_store.acquire_jobs(self.identity, available_slots)
                    for job in jobs:
                        task = self.data_store.get_task(job.task_id)
                        self._running_jobs.add(job.id)
                        executor.submit(copy_context().run, self._run_job, job, task.func)

                self._wakeup_event.wait()
                self._wakeup_event = threading.Event()
        except BaseException as exc:
            self.logger.exception('Worker crashed')
            exception = exc
        else:
            self.logger.info('Worker stopped')
        finally:
            current_worker.reset(token)
            self._state = RunState.stopped

            self._events.publish(WorkerStopped(exception=exception))