import logging from datetime import datetime import six from apscheduler.job import Job from apscheduler.util import undefined from apscheduler.schedulers.base import STATE_STOPPED from apscheduler.schedulers.asyncio import AsyncIOScheduler from apscheduler.executors.asyncio import AsyncIOExecutor from apscheduler.events import EVENT_JOB_EXECUTED, EVENT_JOB_MAX_INSTANCES, EVENT_JOB_ERROR, EVENT_JOB_MISSED class AsyncIOPriorityScheduler(AsyncIOScheduler): def add_job(self, func, trigger=None, args=None, kwargs=None, id=None, name=None, misfire_grace_time=undefined, coalesce=undefined, max_instances=undefined, next_run_time=undefined, jobstore='default', executor='default', replace_existing=False, group=0, priority=0, **trigger_args): job_kwargs = { 'trigger': self._create_trigger(trigger, trigger_args), 'executor': executor, 'func': func, 'args': tuple(args) if args is not None else (), 'kwargs': dict(kwargs) if kwargs is not None else {}, 'id': id, 'name': name, 'misfire_grace_time': misfire_grace_time, 'coalesce': coalesce, 'max_instances': max_instances, 'next_run_time': next_run_time, 'group': group, 'priority': priority } job_kwargs = dict((key, value) for key, value in job_kwargs.items() if value is not undefined) job = _PriorityJob(self, **job_kwargs) # Don't really add jobs to job stores before the scheduler is up and running with self._jobstores_lock: if self.state == STATE_STOPPED: self._pending_jobs.append((job, jobstore, replace_existing)) self._logger.info('Adding job tentatively -- it will be properly scheduled when ' 'the scheduler starts') else: self._real_add_job(job, jobstore, replace_existing) return job def get_due_jobs(self, now=None): if now is None: now = datetime.now(self.timezone) with self._jobstores_lock: jobs = [] for alias, store in six.iteritems(self._jobstores): jobs.extend(store.get_due_jobs(now)) return jobs def _process_jobs(self): due_jobs = self.get_due_jobs() wait_seconds = super()._process_jobs() executors = [] for job in due_jobs: executor = self._lookup_executor(job.executor) if not executor in executors: executors.append(executor) for executor in executors: if hasattr(executor, 'commit'): executor.commit() return wait_seconds class _PriorityJob(Job): def __init__(self, scheduler, priority=1, group=1, **kwargs): self.priority = priority self.group = group super().__init__(scheduler, **kwargs) def modify(self, **changes): if 'priority' in changes: self.priority = changes.pop('priority') if 'group' in changes: self.group = changes.pop('group') super().modify(self, **changes) def __getstate__(self): state = super().__getstate__() state['priority'] = self.priority state['group'] = self.group return state def __setstate__(self, state): priority = state.pop('priority') group = state.pop('group') super().__setstate__(state) self.priority = priority self.group = group class AsyncIOPriorityExecutor(AsyncIOExecutor): def __init__(self, *args, **kwargs): self.scheduler = None self.job_groups = {} super().__init__(*args, **kwargs) def start(self, scheduler, alias): # Store a reference to the scheduler and add job listener self.scheduler = scheduler self.scheduler.add_listener(self.on_job_executed, mask=EVENT_JOB_EXECUTED|EVENT_JOB_MAX_INSTANCES|EVENT_JOB_ERROR|EVENT_JOB_MISSED) super().start(scheduler, alias) def shutdown(self, *args): # Remove job listener self.scheduler.remove_listener(self.on_job_executed) super().shutdown(*args) def submit_job(self, job, run_times): run_times = set(run_times) logger = logging.getLogger('zasd') logger.debug('Incoming submission for job %s with run times:', job.id) for time in run_times: logger.debug(' %s', time) if job.group in self.job_groups: job_group = self.job_groups[job.group] if job.id in job_group['times_for_job']: # Fetch time set times_for_job = job_group['times_for_job'][job.id] else: # Create and store time set times_for_job = set() job_group['times_for_job'][job.id] = times_for_job # Fetch map of times to their jobs job_for_time = job_group['job_for_time'] # Filter out run times that coincide with higher-priority jobs run_times = set( time for time in run_times if not time in job_for_time or job.priority < self.scheduler.get_job(job_for_time[time]).priority ) else: # Create and store empty job group times_for_job = set() job_for_time = {} job_group = { 'times_for_job': { job.id: times_for_job }, 'job_for_time': job_for_time } self.job_groups[job.group] = job_group # Add new times to stored time set for current job times_for_job |= run_times # Remove jobs in time set from other jobs in group # Necessary when incoming job overwrites time slots for existing jobs # # Look at time sets for all jobs in group for job_id, times in list(job_group['times_for_job'].items()): # Look at time set for this job for time in set(times): # Does time in set coincide with time set of job being submitted? if time in times_for_job and job_id != job.id: # Remove time from time set times.remove(time) if len(times) == 0: # Delete time set if empty del job_group['times_for_job'][job_id] else: # Update time set if not job_group['times_for_job'][job_id] = times logger.debug('Final time set for job %s:', job.id) for time in times_for_job: logger.debug(' %s', time) # Map run times to jobs for time in run_times: job_for_time[time] = job.id def commit(self): logger = logging.getLogger('zasd') logger.debug('Committing jobs:') # Look at every group for group_id, group in self.job_groups.items(): # Look at every job in group for job_id, times in group['times_for_job'].items(): # Find job for ID and submit to scheduler superclass job = self.scheduler.get_job(job_id) super().submit_job(job, list(times)) logger.debug(' %s', job.id) for time in times: logger.debug(' %s', time) def on_job_executed(self, event): time = event.scheduled_run_time logger = logging.getLogger('zasd') logger.debug('Job %s has executed', event.job_id) # Find job for job ID job = self.scheduler.get_job(event.job_id) # Get group and times for job job_group = self.job_groups[job.group] times_for_job = job_group['times_for_job'] job_for_time = job_group['job_for_time'] # Delete job from time map del job_for_time[time] # Delte time from job time set times_for_job[job.id].remove(time) if len(times_for_job[job.id]) == 0: logger.debug('Deleting empty time set for job %s', job.id) del times_for_job[job.id] else: logger.debug('Remaining time set for job %s:', job.id) for time in times_for_job[job.id]: logger.debug(' %s', time)