elstat/elstat/adapters.py

import asyncio
import time
import re

PING_RGX = re.compile(r'(.+)( 0% packet loss)(.+)', re.I | re.M)


class Adapter:
    spec = {
        'db': None,
    }

    @classmethod
    async def query(cls, _worker, _adp_args) -> tuple:
        """Main query function."""
        raise NotImplementedError


class PingAdapter(Adapter):
    """Ping the given address and report if
    any packet loss happened."""
    spec = {
        'db': ('timestamp', 'status')
    }

    @classmethod
    async def query(cls, worker, adp_args: dict):
        process = await asyncio.create_subprocess_shell(
            f'ping -c 1 {adp_args["address"]}',
            stderr=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
        )

        out, err = map(lambda s: s.decode('utf-8'),
                       await process.communicate())

        alive = bool(re.search(PING_RGX, out + err))
        worker.log.info(f'{worker.name}: alive? {alive}')

        return (alive,)


class HttpAdapter(Adapter):
    """Adapter to check if a certain
    URL is giving 200."""
    spec = {
        'db': ('timestamp', 'status', 'latency')
    }

    @classmethod
    async def query(cls, worker, adp_args: dict):
        # yes, lots of attributes
        session = worker.manager.app.session

        t_start = time.monotonic()
        resp = await session.get(f'{adp_args["url"]}')
        t_end = time.monotonic()

        latency = round((t_end - t_start) * 1000)

        worker.log.info(f'{worker.name}: status={resp.status} '
                        f'latency={latency}ms')

        if resp.status == 200:
            return True, latency

        # use 0ms drops as failures
        return False, 0