Job scheduler bugfixes and refactoring (#351)

* Job scheduler bugfixes and refactoring

* Cleanup

* Add missing typehints

Author: droserasprout

CHANGELOG.md

@@ -6,6 +6,11 @@
 
 * tzkt: Added `originated_contract_tzips` field to `OperationData`.
 
+### Fixed
+
+* jobs: Fixed jobs with `daemon` schedule never start.
+* jobs: Fixed failed jobs not throwing exceptions into the main loop.
+
 ## 5.1.1 - 2022-05-13
 
 ### Fixed

pyproject.toml

@@ -90,7 +90,7 @@ extended_default_ignore = []
 
 [tool.flakehell.plugins]
 pyflakes = ["+*"]
-"flake8-*" = ["+*"]
+"flake8-*" = ["+*", "-C417"]
 flake8-docstrings = ["-*"]
 flake8-simplify = ["-SIM106"]
 

src/dipdup/context.py

@@ -382,6 +382,7 @@ def __init__(self, package: str) -> None:
         self._hooks: Dict[str, HookConfig] = {}
 
     async def run(self) -> None:
+        self._logger.debug('Starting CallbackManager loop')
         while True:
             async with slowdown(1):
                 while pending_hooks:

src/dipdup/dipdup.py

@@ -7,7 +7,6 @@
 from asyncio import gather
 from collections import deque
 from contextlib import AsyncExitStack
-from contextlib import asynccontextmanager
 from contextlib import suppress
 from typing import Awaitable
 from typing import Deque
@@ -16,7 +15,6 @@
 from typing import Set
 from typing import Tuple
 
-from apscheduler.events import EVENT_JOB_ERROR  # type: ignore
 from prometheus_client import start_http_server  # type: ignore
 from tortoise.exceptions import OperationalError
 from tortoise.transactions import get_connection
@@ -59,8 +57,7 @@
 from dipdup.models import OperationData
 from dipdup.models import Schema
 from dipdup.prometheus import Metrics
-from dipdup.scheduler import add_job
-from dipdup.scheduler import create_scheduler
+from dipdup.scheduler import SchedulerManager
 from dipdup.utils import is_importable
 from dipdup.utils import slowdown
 from dipdup.utils.database import generate_schema
@@ -367,7 +364,6 @@ def __init__(self, config: DipDupConfig) -> None:
             datasources=self._datasources,
             callbacks=self._callbacks,
         )
-        self._scheduler = create_scheduler(self._config.advanced.scheduler)
         self._codegen = DipDupCodeGenerator(self._config, self._datasources_by_config)
         self._schema: Optional[Schema] = None
 
@@ -492,19 +488,19 @@ async def _set_up_database(self, stack: AsyncExitStack) -> None:
     async def _set_up_hooks(self, tasks: Optional[Set[Task]] = None) -> None:
         for hook_config in default_hooks.values():
             try:
-                self._ctx.callbacks.register_hook(hook_config)
+                self._callbacks.register_hook(hook_config)
             except ProjectImportError:
                 if hook_config.callback in ('on_rollback', 'on_index_rollback'):
-                    self._logger.info(f'Hook {hook_config.callback} is not available')
+                    self._logger.info(f'Hook `{hook_config.callback}` is not available')
                 else:
                     raise
 
         for hook_config in self._config.hooks.values():
-            self._ctx.callbacks.register_hook(hook_config)
+            self._callbacks.register_hook(hook_config)
 
         # FIXME: Why does `is not None` check break oneshot mode?
         if tasks:
-            tasks.add(create_task(self._ctx.callbacks.run()))
+            tasks.add(create_task(self._callbacks.run()))
 
     async def _set_up_prometheus(self) -> None:
         if self._config.prometheus:
@@ -596,41 +592,14 @@ async def _event_wrapper():
         return event  # noqa: R504
 
     async def _set_up_scheduler(self, stack: AsyncExitStack, tasks: Set[Task]) -> Event:
-        job_failed = Event()
+        # NOTE: Prepare SchedulerManager
         event = Event()
-        exception: Optional[Exception] = None
+        scheduler = SchedulerManager(self._config.advanced.scheduler)
+        run_task = create_task(scheduler.run(event))
+        tasks.add(run_task)
 
-        @asynccontextmanager
-        async def _context():
-            try:
-                self._scheduler.start()
-                yield
-            finally:
-                self._scheduler.shutdown()
-
-        def _error_hook(event) -> None:
-            nonlocal job_failed, exception
-            exception = event.exception
-            job_failed.set()
-
-        async def _watchdog() -> None:
-            nonlocal job_failed
-            await job_failed.wait()
-            if not isinstance(exception, Exception):
-                raise RuntimeError
-            raise exception
-
-        async def _event_wrapper():
-            self._logger.info('Waiting for an event to start scheduler')
-            await event.wait()
+        # NOTE: Register jobs
+        for job_config in self._config.jobs.values():
+            scheduler.add_job(self._ctx, job_config)
 
-            self._logger.info('Starting scheduler')
-            self._scheduler.add_listener(_error_hook, EVENT_JOB_ERROR)
-            await stack.enter_async_context(_context())
-            tasks.add(create_task(_watchdog()))
-
-            for job_config in self._config.jobs.values():
-                add_job(self._ctx, self._scheduler, job_config)
-
-        tasks.add(create_task(_event_wrapper()))
         return event  # noqa: R504

src/dipdup/scheduler.py

@@ -1,15 +1,19 @@
 import asyncio
 import json
-from datetime import datetime
+import logging
 from functools import partial
 from typing import Any
 from typing import Dict
 from typing import Optional
+from typing import Set
 
+from apscheduler.events import EVENT_JOB_ERROR  # type: ignore
+from apscheduler.events import EVENT_JOB_EXECUTED
+from apscheduler.events import JobEvent  # type: ignore
+from apscheduler.job import Job  # type: ignore
 from apscheduler.schedulers.asyncio import AsyncIOScheduler  # type: ignore
 from apscheduler.triggers.cron import CronTrigger  # type: ignore
 from apscheduler.triggers.interval import IntervalTrigger  # type: ignore
-from apscheduler.util import undefined  # type: ignore
 
 from dipdup.config import JobConfig
 from dipdup.context import DipDupContext
@@ -24,57 +28,95 @@
 }
 
 
-def create_scheduler(config: Optional[Dict[str, Any]] = None) -> AsyncIOScheduler:
-    if not config:
-        return AsyncIOScheduler(DEFAULT_CONFIG)
-
+def _verify_config(config: Dict[str, Any]) -> None:
+    """Ensure that dict is a valid `apscheduler` config"""
     json_config = json.dumps(config)
     if 'apscheduler.executors.pool' in json_config:
         raise ConfigurationError('`apscheduler.executors.pool` is not supported. If needed, create a pool inside a regular hook.')
     for key in config:
         if not key.startswith('apscheduler.'):
             raise ConfigurationError('`advanced.scheduler` config keys must start with `apscheduler.`, see apscheduler library docs')
-    return AsyncIOScheduler(config)
 
 
-def add_job(ctx: DipDupContext, scheduler: AsyncIOScheduler, job_config: JobConfig) -> None:
-    hook_config = job_config.hook_config
+class SchedulerManager:
+    def __init__(self, config: Optional[Dict[str, Any]] = None) -> None:
+        if config:
+            _verify_config(config)
+        self._logger = logging.getLogger('dipdup.jobs')
+        self._scheduler = AsyncIOScheduler(config or DEFAULT_CONFIG)
+        self._scheduler.add_listener(self._on_error, EVENT_JOB_ERROR)
+        self._scheduler.add_listener(self._on_executed, EVENT_JOB_EXECUTED)
+        self._exception: Optional[Exception] = None
+        self._exception_event: asyncio.Event = asyncio.Event()
+        self._daemons: Set[str] = set()
 
-    async def _job_wrapper(ctx: DipDupContext, *args, **kwargs) -> None:
-        nonlocal job_config, hook_config
-        job_ctx = HookContext(
-            config=ctx.config,
-            datasources=ctx.datasources,
-            callbacks=ctx.callbacks,
-            logger=logger,
-            hook_config=hook_config,
-        )
+    async def run(self, event: asyncio.Event) -> None:
+        self._logger.info('Waiting for an event to start scheduler')
+        await event.wait()
+
+        self._logger.info('Starting scheduler')
         try:
-            await job_ctx.fire_hook(hook_config.callback, *args, **kwargs)
+            self._scheduler.start()
+            await self._exception_event.wait()
+            if self._exception is None:
+                raise RuntimeError('Job has failed but exception is not set')
+            raise self._exception
         except asyncio.CancelledError:
             pass
+        finally:
+            self._scheduler.shutdown()
+
+    def add_job(self, ctx: DipDupContext, job_config: JobConfig) -> Job:
+        if job_config.daemon:
+            self._daemons.add(job_config.name)
+
+        hook_config = job_config.hook_config
+
+        logger = FormattedLogger(
+            name=f'dipdup.hooks.{hook_config.callback}',
+            fmt=job_config.name + ': {}',
+        )
+
+        async def _job_wrapper(ctx: DipDupContext, *args, **kwargs) -> None:
+            nonlocal job_config, hook_config
+            job_ctx = HookContext(
+                config=ctx.config,
+                datasources=ctx.datasources,
+                callbacks=ctx.callbacks,
+                logger=logger,
+                hook_config=hook_config,
+            )
+            try:
+                await job_ctx.fire_hook(
+                    hook_config.callback,
+                    *args,
+                    **kwargs,
+                )
+            except asyncio.CancelledError:
+                pass
+
+        if job_config.crontab:
+            trigger = CronTrigger.from_crontab(job_config.crontab)
+        elif job_config.interval:
+            trigger = IntervalTrigger(seconds=job_config.interval)
+        elif job_config.daemon:
+            trigger = None
         else:
-            if job_config.daemon:
-                raise ConfigurationError('Daemon jobs are intended to run forever')
-
-    logger = FormattedLogger(
-        name=f'dipdup.hooks.{hook_config.callback}',
-        fmt=job_config.name + ': {}',
-    )
-    if job_config.crontab:
-        trigger, next_run_time = CronTrigger.from_crontab(job_config.crontab), undefined
-    elif job_config.interval:
-        trigger, next_run_time = IntervalTrigger(seconds=job_config.interval), undefined
-    elif job_config.daemon:
-        trigger, next_run_time = None, datetime.now()
-    else:
-        raise RuntimeError
-
-    scheduler.add_job(
-        func=partial(_job_wrapper, ctx=ctx),
-        id=job_config.name,
-        name=job_config.name,
-        trigger=trigger,
-        next_run_time=next_run_time,
-        kwargs=job_config.args,
-    )
+            raise RuntimeError
+
+        return self._scheduler.add_job(
+            func=partial(_job_wrapper, ctx=ctx),
+            id=job_config.name,
+            name=job_config.name,
+            trigger=trigger,
+            kwargs=job_config.args,
+        )
+
+    def _on_error(self, event: JobEvent) -> None:
+        self._exception = event.exception
+        self._exception_event.set()
+
+    def _on_executed(self, event: JobEvent) -> None:
+        if event.job_id in self._daemons:
+            event.exception = ConfigurationError('Daemon jobs are intended to run forever')
+            self._on_error(event)