redis锁

# -*- coding: utf-8 -*-

from contextlib import contextmanager
import redis
from config import REDIS_HOST, REDIS_PORT, REDIS_PASS, logger


TE_BALANCE_LOCK_CLI = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=9, password=REDIS_PASS)


class TELock(object):
    """ redis做的锁 """

    def __init__(
        self,
        account_no, # 作为锁的唯一条件
        max_retries=5,  # 最大重试次数
        db_ex_time=10   # 过期时间
    ):
        self.cli = TE_BALANCE_LOCK_CLI
        self.account_no = account_no
        self.max_retries = max_retries
        self.db_ex_time = db_ex_time

    def __get_lock(self):
        delay = 1   # 重试间隔时间
        times = 0   # 重试次数
        while True:
            # 获取为1或者获取不到都视为没有锁
            if self.cli.set(
                self.account_no,
                "1",
                ex=self.db_ex_time, # 过期时间
                nx=True # nx设置为True，则只有self.account_no不存在时，当前set操作才执行
            ):
                return True
            logger.info(
                "{} is in use, {}s 后重试".format(self.account_no, delay)
            )
            times += 1
            if times > self.max_retries:
                raise Exception("wait {} balance lock too long.".format(self.account_no))
            time.sleep(delay)

    @property
    @contextmanager
    def balance_lock(self):
        self.__get_lock()
        try:
            yield self
        finally:
            self.cli.delete(self.account_no)


if __name__ == '__main__':
    import time

    with TELock("9999").balance_lock:
        print "get"
        time.sleep(1000)

        print "exit"