Celery `redbeat` scheduler compatability about celery-pool-asyncio HOT 12 OPEN

from celery import beat
from redbeat import RedBeatScheduler
from redbeat.schedulers import get_redis, logger


class AsyncSchedulerMixin:
    async def apply_async(self, entry, producer=None, advance=True, **kwargs):
        # Update time-stamps and run counts before we actually execute,
        # so we have that done if an exception is raised (doesn't schedule
        # forever.)
        entry = self.reserve(entry) if advance else entry
        task = self.app.tasks.get(entry.task)

        try:
            entry_args = [
                v() if isinstance(v, beat.BeatLazyFunc) else v
                for v in (entry.args or [])
            ]
            entry_kwargs = {
                k: v() if isinstance(v, beat.BeatLazyFunc) else v
                for k, v in entry.kwargs.items()
            }

            if task:
                return await task.apply_async(
                    entry_args, entry_kwargs,
                    producer=producer,
                    **entry.options,
                )
            else:
                return await self.send_task(
                    entry.task, entry_args, entry_kwargs,
                    producer=producer,
                    **entry.options,
                )
        except Exception as exc:  # pylint: disable=broad-except
            msg = "Couldn't apply scheduled task {entry.name}: {exc}".format(
                entry=entry,
                exc=exc,
            )
            beat.reraise(
                beat.SchedulingError,
                beat.SchedulingError(msg),
                beat.sys.exc_info()[2]
            )
        finally:
            self._tasks_since_sync += 1
            if self.should_sync():
                self._do_sync()

    async def apply_entry(self, entry, producer=None):
        beat.info(
            'Scheduler: Sending due task %s (%s)',
            entry.name,
            entry.task,
        )

        try:
            coro = self.apply_async(
                entry=entry,
                producer=producer,
                advance=False,
            )
            result = await coro
        except Exception as exc:  # pylint: disable=broad-except
            beat.error(
                'Message Error: %s\n%s',
                exc,
                beat.traceback.format_stack(),
                exc_info=True,
            )
        else:
            beat.debug(
                '%s sent. id->%s',
                entry.task,
                result.id,
            )


class RedBeatAsyncScheduler(AsyncSchedulerMixin, RedBeatScheduler):

    async def maybe_due(self, entry, **kwargs):
        is_due, next_time_to_run = entry.is_due()

        if is_due:
            logger.info('Scheduler: Sending due task %s (%s)', entry.name, entry.task)
            try:
                result = await self.apply_async(entry, **kwargs)
            except Exception as exc:
                logger.exception('Message Error: %s', exc)
            else:
                logger.debug('%s sent. id->%s', entry.task, result.id)
        return next_time_to_run

    async def tick(self, min=min, **kwargs):
        if self.lock:
            logger.debug('beat: Extending lock...')
            get_redis(self.app).pexpire(self.lock_key, int(self.lock_timeout * 1000))

        remaining_times = []
        try:
            for entry in self.schedule.values():
                next_time_to_run = await self.maybe_due(entry, **self._maybe_due_kwargs)
                if next_time_to_run:
                    remaining_times.append(next_time_to_run)
        except RuntimeError:
            logger.debug('beat: RuntimeError', exc_info=True)

        return min(remaining_times + [self.max_interval])