from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta
from logging import getLogger
from multiprocessing.queues import Queue
from typing import Optional, cast
from edri.abstract import ManagerBase
from edri.config.constant import SCHEDULER_TIMEOUT_MAX
from edri.config.setting import HEALTH_CHECK_LOG
from edri.dataclass.event import Event
from edri.dataclass.response import ResponseStatus
from edri.events.edri.router import HealthCheck
from edri.events.edri.scheduler import Cancel, Set, SetUpdate, Update
from edri.utility import Storage
__all__ = ["Job"]
[docs]
@dataclass
class Job:
"""
Represents a scheduled task within the scheduler.
The Job dataclass intentionally does not normalize its own values. Time
normalization is handled by Scheduler when accepting external input, so
Job remains a simple internal data object.
Attributes:
event (Event): The event to be triggered when the job is executed.
when (datetime): The timezone-aware time at which the job should be
executed.
repeat (Optional[timedelta]): How often the job should repeat. If None,
the job does not repeat.
"""
event: Event
when: datetime
repeat: Optional[timedelta] = None
class Scheduler(ManagerBase):
"""
A scheduler for managing and executing timed tasks.
The scheduler accepts scheduling events, stores jobs, executes due jobs,
and optionally reschedules repeating jobs.
"""
def __init__(self, router_queue: "Queue[Event]", jobs: list[Job] | Mapping[str, Job]) -> None:
"""
Initializes the Scheduler instance.
Initial jobs may be given either as a list, in which case the scheduler
generates an identifier for each job, or as a mapping, in which case the
mapping keys become the job identifiers and can be used by Update and
Cancel events.
Args:
router_queue (Queue): Queue used for routing events.
jobs (list[Job] | Mapping[str, Job]): Initial jobs to be loaded when
the scheduler starts.
"""
super().__init__(router_queue, getLogger(__name__))
self.jobs: Storage[Job]
self.initial_jobs: list[tuple[Optional[str], Job]]
if isinstance(jobs, Mapping):
self.initial_jobs = list(jobs.items())
else:
self.initial_jobs = [(None, job) for job in jobs]
@staticmethod
def _normalize_datetime(value: datetime) -> datetime:
"""
Normalize datetime values accepted by the scheduler.
Naive datetime values are interpreted as system-local time and converted
to timezone-aware datetime instances. Timezone-aware values are returned
unchanged, because Python can compare aware datetimes across timezones.
Args:
value (datetime): Datetime value to normalize.
Returns:
datetime: Timezone-aware datetime.
"""
if value.tzinfo is None or value.utcoffset() is None:
return value.astimezone()
return value
def after_start(self) -> None:
"""
Initializes job storage and loads initial jobs.
Initial jobs are treated as scheduler input, so their execution time is
normalized before they are stored. Jobs that came with an identifier are
stored under it; the rest get a generated one.
"""
super().after_start()
self.jobs = Storage[Job]()
for key, job in self.initial_jobs:
job.when = self._normalize_datetime(job.when)
try:
self.jobs.append(job, key)
except KeyError as e:
self.logger.warning("Scheduled task already exists: %s", key, exc_info=e)
def _create_job(self, event: Set | SetUpdate) -> Optional[str]:
"""
Creates and stores a new job from a scheduler event.
A new job requires both the event to execute and the first execution
time. This matters especially for SetUpdate, where some fields may be
optional when updating an existing job, but are still required when a
new job has to be created.
Args:
event (Set | SetUpdate): Scheduler event containing job data.
Returns:
Optional[str]: Identifier of the created job, or None if the job
could not be created because required data was missing.
"""
scheduled_event = event.event
when = event.when
if scheduled_event is None or when is None:
event.response.set_status(ResponseStatus.FAILED)
self.logger.debug(
"Scheduled task %s cannot be created without event and when",
event.identifier,
)
return None
key = self.jobs.append(
Job(
scheduled_event,
self._normalize_datetime(when),
event.repeat,
),
event.identifier,
)
if not event.identifier:
event.response.identifier = key
self.logger.debug("New scheduled task: %s", key)
return key
def _update_job_from_event(self, job: Job, event: Update | SetUpdate) -> None:
"""
Updates an existing job from a scheduler event.
Only fields explicitly provided by the event are updated. If event.when
is None, the current execution time remains unchanged.
Args:
job (Job): Job instance to update.
event (Update | SetUpdate): Event containing updated job values.
"""
if event.event is not None:
job.event = event.event
if event.when is not None:
job.when = self._normalize_datetime(event.when)
if event.repeat is not None:
job.repeat = event.repeat
def solve_req_set(self, event: Set) -> None:
"""
Handles a request to schedule a new job.
Args:
event (Set): Event containing the job to schedule.
"""
try:
self._create_job(event)
except KeyError as e:
self.logger.warning(
"Scheduled task already exists: %s",
event.identifier,
exc_info=e,
)
def solve_req_update(self, event: Update) -> None:
"""
Handles a request to update an existing job.
Args:
event (Update): Event containing the job identifier and fields to
update.
"""
job = self.jobs.get(event.identifier, None)
if not job:
event.response.set_status(ResponseStatus.FAILED)
self.logger.debug("Scheduled task %s was not found", event.identifier)
return
self._update_job_from_event(job, event)
self.logger.debug("Scheduled task %s was updated", event.identifier)
def solve_req_set_update(self, event: SetUpdate) -> None:
"""
Handles a request to create or update a job.
If the job already exists, only provided fields are updated. If the job
does not exist, both event.event and event.when are required to create
a valid new job.
Args:
event (SetUpdate): Event containing create-or-update data.
"""
job = self.jobs.get(cast(str, event.identifier), None)
if not job:
self._create_job(event)
return
self._update_job_from_event(job, event)
self.logger.debug("Scheduled task %s was updated", event.identifier)
def solve_req_cancel(self, event: Cancel) -> None:
"""
Handles a request to cancel an existing job.
Args:
event (Cancel): Event containing the identifier of the job to cancel.
"""
try:
del self.jobs[event.identifier]
except KeyError:
self.logger.debug("Scheduled task %s was not found", event.identifier)
return
self.logger.debug("Scheduled task %s was cancelled", event.identifier)
def get_next_job(self) -> tuple[str, Optional[Job]]:
"""
Returns the nearest scheduled job.
Returns:
tuple[str, Optional[Job]]: Identifier and job instance of the next
scheduled job. If no job exists, returns an empty identifier and None.
"""
if not self.jobs:
return "", None
if len(self.jobs) == 1:
return next(iter(self.jobs.items()))
jobs = iter(self.jobs.items())
next_key, next_job = next(jobs)
for key, job in jobs:
if next_job.when > job.when:
next_job = job
next_key = key
return next_key, next_job
def run_pending(self) -> None:
"""
Executes all jobs that are due.
Non-repeating jobs are removed after execution. Repeating jobs are moved
forward by their repeat interval.
"""
now = datetime.now().astimezone()
for key, job in list(self.jobs.items()):
if now >= job.when:
self.send_to_router(job.event)
if job.repeat:
job.when = job.when + job.repeat
else:
self.logger.debug("Scheduled task %s was completed", key)
del self.jobs[key]
def run_resolver(self) -> None:
"""
Runs the scheduler loop.
The loop executes pending jobs, calculates the timeout until the next
scheduled job, and listens for incoming scheduler-management events.
"""
while True:
try:
self.run_pending()
self.logger.debug("Count of scheduled task: %s", len(self.jobs))
key, next_job = self.get_next_job()
if not next_job:
timeout = SCHEDULER_TIMEOUT_MAX
else:
now = datetime.now().astimezone()
timeout = (next_job.when - now).total_seconds()
if timeout > SCHEDULER_TIMEOUT_MAX:
timeout = SCHEDULER_TIMEOUT_MAX
self.logger.debug(
"Next scheduled task %s will be executed after %ss",
key,
timeout,
)
if self.router_pipe.poll(timeout=timeout):
event = self.router_pipe.recv()
if HEALTH_CHECK_LOG or not isinstance(event, HealthCheck):
self.logger.debug("Received event: %s", event)
self.resolve(event)
except KeyboardInterrupt:
return