pytools/common/redis_lock.py

# -*- coding: utf-8 -*-
# @Time    : 2019/9/30 22:03
# @Author  : CRJ
# @File    : utils.py
# @Software: PyCharm
# @Python3.6
import math
import time
import uuid

import redis


# 获取锁（乐观锁）
def acquire_lock_with_timeout(conn, lockname, acquire_timeout=5, lock_timeout=10):
    # 128位随机标识符
    identifier = str(uuid.uuid4())
    lockname = 'lock:' + lockname
    lock_timeout = int(math.ceil(lock_timeout))  # 确保传给exprie是整数

    end = time.time() + acquire_timeout
    while time.time() < end:
        if conn.setnx(lockname, identifier):
            conn.expire(lockname, lock_timeout)
            return identifier
        elif not conn.ttl(lockname):  # 为没有设置超时时间的锁设置超时时间
            conn.expire(lockname, lock_timeout)

        time.sleep(0.001)
    return False


# 释放锁（乐观锁）
def release_lock(conn, lockname, identifier):
    lockname = 'lock:' + lockname
    with conn.pipeline(True) as pipe:
        while True:
            try:
                pipe.watch(lockname)
                value = pipe.get(lockname)
                # 判断标志是否相同
                if value is not None and value == identifier:
                    pipe.multi()
                    pipe.delete(lockname)
                    pipe.execute()
                    return True

                # 不同则直接退出 return False
                pipe.unwatch()
                break
            except redis.exceptions.WatchError:
                pass

        return False


class OptimisticLock:

    def __init__(self, conn, lockname, acquire_timeout=60, lock_timeout=10):
        self.conn = conn
        self.lockname = 'lock:' + lockname
        self.acquire_timeout = acquire_timeout
        self.lock_timeout = int(math.ceil(lock_timeout))  # 确保传给 expire 是整数
        self.identifier = None

    def __enter__(self):
        # 128位随机标识符
        self.identifier = str(uuid.uuid4())

        end = time.time() + self.acquire_timeout
        while time.time() < end:
            if self.conn.setnx(self.lockname, self.identifier):
                self.conn.expire(self.lockname, self.lock_timeout)
                return self

            elif not self.conn.ttl(self.lockname):  # 为没有设置超时时间的锁设置超时时间
                self.conn.expire(self.lockname, self.lock_timeout)

            time.sleep(0.001)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        with self.conn.pipeline(True) as pipe:
            while True:
                try:
                    pipe.watch(self.lockname)
                    value = pipe.get(self.lockname)
                    # 判断标志是否相同
                    if value is not None and value == self.identifier:
                        pipe.multi()
                        pipe.delete(self.lockname)
                        pipe.execute()
                        return True

                    # 不同则直接退出 return False
                    pipe.unwatch()
                    break
                except redis.exceptions.WatchError:
                    pass

            return False