Simple multiprocess log re-player for elliptics proxy

replay_logs.py

#!/usr/bin/env python3

#
# Simple multiprocess log re-player for elliptics proxy
#

from hashlib import sha256
from time import sleep
from urllib import request
from urllib.error import HTTPError
from argparse import ArgumentParser
from itertools import islice, repeat
from functools import wraps
from concurrent.futures import ProcessPoolExecutor

verbose = False
def debug(*args, **kwargs):
    """
    Prints debug message if verbose mode is on

    TODO: rewrite using logging
    """
    if verbose:
        print(*args, **kwargs)

def retry_on_exception(retries=3, base=2.0, backoff='exponential'):
    backoffs = {
            'exponential': lambda: (base**x for x in range(retries)),
            'linear': lambda: (base*x for x in range(retries)),
            'fixed': lambda: islice(repeat(base), retries),
        }
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            exception = RuntimeError
            for wait in backoffs[backoff]():
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    exception = e
                sleep(wait)
            raise exception
        return wrapper
    return decorator

@retry_on_exception()
def _fetch(url):
    """
    Low-level fetch function
    """
    return request.urlopen(url)

def fetch(url):
    """
    Tries to fetch an URL.
    """
    exc = None
    size = 0
    checksum = ""
    try:
        response = _fetch(url)
        size = response.length
        checksum = sha256(response.readall()).hexdigest()
    except HTTPError as e:
        exc = "{}: {}".format(repr(e), e.msg)
    except Exception as e:
        exc = repr(e)
    return {'url': url, 'checksum': checksum, 'size': int(size), 'exception': exc}


if __name__ == '__main__':
    parser = ArgumentParser(description='Replay urls against proxy.')
    parser.add_argument('files', metavar='FILE', type=str, nargs='+',
        help='files for replay')
    parser.add_argument('--url', dest='url', action='store', default="",
        help='url to prepend to entries')
    parser.add_argument('--threads', dest='threads', type=int, action='store',
        default=4, help='number of threads')
    parser.add_argument('--debug', dest='debug', action='store_true',
        help='debug output')
    args = parser.parse_args()

    verbose = args.debug

    total = 0
    failed = 0
    for file in args.files:
        with open(file) as lines:
            with ProcessPoolExecutor(max_workers=args.threads) as executor:
                results = executor.map(fetch, (args.url + line.strip() for line in lines))
                for result in results:
                    if result['exception']:
                        debug("FAILED:\t{url}:\t{exception}".format(**result))
                        failed += 1
                    else:
                        debug("OK:\t{url}:\t{size}\t{checksum}".format(**result))
                    total += 1

    print("Total:\t{}".format(total))
    print("Failed:\t{}".format(failed))