redis_latency.py

import logging
import os
import time

import redis
import statsd

logger = logging.getLogger(__name__)


class Timer(object):
    def __enter__(self):
        self.start = time.time()
        return self

    def __exit__(self, *args):
        self.end = time.time()
        self.interval = self.end - self.start


class RedisLatencyCollector(object):
    def __init__(self, host, port, auth=None):
        self.connection = None
        self.host = host
        self.port = port
        self.auth = auth

        # Warm the connection
        self._client(host, port, auth)

    def _init_connection(self, host, port, db, auth, timeout):
        """Return a redis connection, cached if possible"""
        if self.connection:
            return self.connection

        try:
            cli = redis.Redis(host=host, port=port,
                              db=db, socket_timeout=timeout, password=auth)
        except Exception, ex:
            logger.error("RedisLatencyCollector: failed to connect to {}:{}. {}.".format(host, port, ex))
            return None

        self.connection = cli

        return cli

    def _client(self, host, port, auth=None, timeout=5):
        """Return a redis client for the configuration."""

        db = 0
        timeout = timeout

        try:
            cli = self._init_connection(host, port, db, auth, timeout=timeout)
            cli.ping()
            return cli
        except Exception, ex:
            logger.error("RedisLatencyCollector: failed to ping {}:{}. {}.".format(host, port, ex))
            self.connection = None
            return None

    def collect_latency(self):
        client = self._client(self.host, self.port, self.auth)

        with Timer() as t:
            client.ping()
        return t.interval * 1000


def main():
    host = os.environ['REDIS_HOST']
    port = os.environ.get('REDIS_PORT', 6379)
    auth = os.environ.get('REDIS_AUTH', None)
    interval = int(os.environ.get('POLL_INTERVAL', '250'))
    statsd_host = os.environ['STATSD_HOST']
    statsd_port = int(os.environ.get('STATSD_PORT', 8125))
    statsd_prefix = os.environ.get('STATSD_PREFIX', 'test')

    logger.debug("Redis config {}:{}".format(host, port))
    logger.debug("Statsd config {}:{} prefix - {}".format(statsd_host, statsd_port, statsd_prefix))

    collector = RedisLatencyCollector(host, port, auth)

    statsd_client = statsd.StatsClient(statsd_host, statsd_port, prefix=statsd_prefix)
    statsd_key = 'redis.{}.latency'.format(host.split('.')[0])

    while True:
        latency = collector.collect_latency()
        statsd_client.timing(statsd_key, latency)
        time.sleep(float(interval) / 1000)


if __name__ == '__main__':
    main()