diff --git a/isitdown/checks/base.py b/isitdown/checks/base.py index 29d64bb..e61aa28 100644 --- a/isitdown/checks/base.py +++ b/isitdown/checks/base.py @@ -2,6 +2,8 @@ import logging import asyncio +from enum import Enum + from ..notifiers import LoggingNotifier LOG = logging.getLogger(__name__) @@ -15,43 +17,44 @@ class Result: reason = attr.ib(default="") +class state(Enum): + STOPPED = 0 + RUNNING = 1 + PAUSED = 2 + + class BaseChecks: def __init__(self, name, *, startup_delay=0, check_interval=300, notifiers=None): - self.name = name - self.running = True - self.check_interval = check_interval - self.startup_delay = startup_delay if not notifiers: notifiers = [ - LoggingNotifier( - logger=logging.getLogger(f"isitdown.status.{self.name}") - ) + LoggingNotifier(logger=logging.getLogger(f"isitdown.status.{name}")) ] + + self.name = name + self._task = None + self.state = state.RUNNING self.notifiers = notifiers + self.startup_delay = startup_delay + self.paused_seconds = 0 + self.check_interval = check_interval - async def start(self): + def start(self): + loop = asyncio.get_event_loop() + self._task = loop.create_task(self._start()) + return self._task + + async def _start(self): LOG.info(f"Starting check: {self.name}") try: - await self.startup() await asyncio.sleep(self.startup_delay) - while self.running: - try: - result = await self.check() - except Exception as e: - LOG.debug(f"Error during check: {self.name}", exc_info=1) - result = Result( - check=self.name, - success=False, - reason=f"python exception: {e}", - data=e, - ) - finally: - if result.success: - await self.ok(result=result) - else: - await self.error(result=result) - await asyncio.sleep(self.check_interval) + await self.startup() + while self.state in (state.RUNNING, state.PAUSED): + if self.state == state.PAUSED: + await self._pause() + else: + await self._run() + await asyncio.sleep(self.check_interval) except asyncio.CancelledError: pass except Exception as e: @@ -59,6 +62,35 @@ async def start(self): finally: await self.shutdown() + async def _run(self): + try: + result = await self.check() + except Exception as e: + LOG.debug(f"Error during check: {self.name}", exc_info=1) + result = Result( + check=self.name, success=False, reason=f"python exception: {e}", data=e + ) + finally: + if result.success: + await self.ok(result=result) + else: + await self.error(result=result) + + async def _pause(self): + if self.paused_seconds <= 0: + self.state = state.RUNNING + else: + await asyncio.sleep(1) + self.paused_seconds -= 1 + + def pause(self, seconds): + self.state = state.PAUSED + self.paused_seconds = int(seconds) + + async def stop(self): + self._task.cancel() + await self._task + async def check(self): LOG.debug(f"Checking with check: {self.name}") @@ -67,7 +99,7 @@ async def startup(self): async def shutdown(self): LOG.info(f"Shutting down check: {self.name}") - self.running = False + self.running = state.STOPPED async def error(self, result): await asyncio.gather(*[notifier.error(result) for notifier in self.notifiers]) diff --git a/isitdown/isitdown.py b/isitdown/isitdown.py index 43cc4db..94b637c 100644 --- a/isitdown/isitdown.py +++ b/isitdown/isitdown.py @@ -1,5 +1,6 @@ import asyncio import logging +import functools import aiohttp.web LOG = logging.getLogger(__name__) @@ -9,23 +10,21 @@ class isitdown(aiohttp.web.Application): def __init__(self, *checks, **kwargs): super().__init__(**kwargs) if checks is None: - checks = dict() + checks = list() - self.checks = checks - self._checks = dict() - - self.on_startup.append(self._start_checks) + self.checks = dict() + self.on_startup.append(functools.partial(self._start_checks, checks)) self.on_shutdown.append(self._cleanup_checks) def start(self, **kwargs): LOG.info("Starting isitdown") aiohttp.web.run_app(self, **kwargs) - async def _start_checks(self, isitdown): - for check in self.checks: - self._checks[check] = self.loop.create_task(check.start()) + async def _start_checks(self, checks, app): + for check in checks: + self.checks[check.name] = check + check.start() - async def _cleanup_checks(self, isitdown): - for task in self._checks.values(): - task.cancel() - await asyncio.gather(*self._checks.values()) + async def _cleanup_checks(self, app): + tasks = [check.stop() for check in self.checks.values()] + await asyncio.gather(*tasks) diff --git a/isitdown/notifiers/__init__.py b/isitdown/notifiers/__init__.py index 78f65e4..85abc2d 100644 --- a/isitdown/notifiers/__init__.py +++ b/isitdown/notifiers/__init__.py @@ -8,7 +8,7 @@ LOG = logging.getLogger(__name__) -class State(Enum): +class state(Enum): OK = 0 ERROR = 1 @@ -16,11 +16,11 @@ class State(Enum): class BaseNotifier: def __init__(self, notify_after=1): self.errors = 0 - self.state = State.OK + self.state = state.OK self.notify_after = notify_after async def error(self, result): - self.state = State.ERROR + self.state = state.ERROR self.errors += 1 if self.errors % self.notify_after == 0: await self._error(result) @@ -28,8 +28,8 @@ async def error(self, result): await self._silenced_error(result) async def ok(self, result): - if self.state == State.ERROR: - self.state = State.OK + if self.state == state.ERROR: + self.state = state.OK await self._recover(result) else: await self._ok(result)