24 февр. 2015 г.

Timeout декоратор

Если функция будет работать долго, и при этом нужно чтобы она работала не дольше определенного количества секунд, можно применить такой вот декоратор:

import errno
import os
import signal
from functools import wraps


class TimeoutError(Exception):
    pass


def timeout(seconds=10, error_message=os.strerror(errno.ETIME)):
    """
    @timeout(60)
    def long_functions():
        <your code here>
    """
    def decorator(func):
        def _handle_timeout(signum, frame):
            raise TimeoutError(error_message)

        def wrapper(*args, **kwargs):
            signal.signal(signal.SIGALRM, _handle_timeout)
            signal.alarm(seconds)
            try:
                result = func(*args, **kwargs)
            finally:
                signal.alarm(0)
            return result

        return wraps(func)(wrapper)

    return decorator

Lock в файл или в redis?

Иногда нужно гарантировать эксклюзивное исполнение функции для определенного набора данных. Например, только одно одновременное изъятие данных из какого-либо источника. На помощь в таких ситуациях приходят локи.

Проще всего создать какой-нибудь lock-файл и проверять его существование. Если есть файл - lock включен.
import os
from contextlib import contextmanager


class LockedException(Exception):
    pass


@contextmanager
def file_lock(lock_file):
    """
    with file_lock('/tmp/file.lock'):
        some code here
    """
    if os.path.exists(lock_file):
        raise LockedException('LockedException')
    else:
        open(lock_file, 'w').write("1")
        try:
            yield
        finally:
            os.remove(lock_file)

Жалко напрягать HDD такими мелкими операциями. Попробуем делать lock в redis:
import redis
from contextlib import contextmanager


r = r = redis.StrictRedis(host='localhost', port=6379, db=0)


class LockedException(Exception):
    pass

@contextmanager
def redis_lock(lock_id):
    """
    with redis_lock('unique_redis_key'):
        some code here
    """
    if r.get(lock_id):
        raise LockedException('LockedException')
    else:
        r.set(lock_id, 1)
        try:
            yield
        finally:
            r.delete(lock_id)

Теперь проверим какой вариант быстрей:
if __name__ == '__main__':
    import time
    from numpy import mean


    count = 10000

    times1 = []
    for i in xrange(0, count):
        start = time.time()
        with file_lock('/tmp/file.lock'):
            pass

        times1.append(float(time.time() - start))

    times2 = []
    for i in xrange(0, count):
        start = time.time()
        with redis_lock('redis_lock'):
            pass

        times2.append(float(time.time() - start))

    print 'file lock:', max(times1), min(times1), mean(times1)
    print 'redis lock:', max(times2), min(times2), mean(times2)

Результаты:

maxminavg
file_lock0.08177495002750.000284910202026 0.000996773791313
redis_lock0.00140619277954 0.000319957733154 0.000346991157532

Lock в redis в 3 раза быстрей.