jumpserver/apps/common/utils/lock.py

from functools import wraps
import threading

from redis_lock import Lock as RedisLock, NotAcquired
from redis import Redis
from django.db import transaction

from common.utils import get_logger
from common.utils.inspect import copy_function_args
from apps.jumpserver.const import CONFIG

logger = get_logger(__file__)


class AcquireFailed(RuntimeError):
    pass


class DistributedLock(RedisLock):
    def __init__(self, name, blocking=True, expire=None, release_lock_on_transaction_commit=False,
                 release_raise_exc=False, auto_renewal_seconds=60*2):
        """
        使用 redis 构造的分布式锁

        :param name:
            锁的名字，要全局唯一
        :param blocking:
            该参数只在锁作为装饰器或者 `with` 时有效。
        :param expire:
            锁的过期时间
        :param release_lock_on_transaction_commit:
            是否在当前事务结束后再释放锁
        :param release_raise_exc:
            释放锁时，如果没有持有锁是否抛异常或静默
        :param auto_renewal_seconds:
            当持有一个无限期锁的时候，刷新锁的时间，具体参考 `redis_lock.Lock#auto_renewal`
        """
        self.kwargs_copy = copy_function_args(self.__init__, locals())
        redis = Redis(host=CONFIG.REDIS_HOST, port=CONFIG.REDIS_PORT, password=CONFIG.REDIS_PASSWORD)

        if expire is None:
            expire = auto_renewal_seconds
            auto_renewal = True
        else:
            auto_renewal = False

        super().__init__(redis_client=redis, name=name, expire=expire, auto_renewal=auto_renewal)
        self._blocking = blocking
        self._release_lock_on_transaction_commit = release_lock_on_transaction_commit
        self._release_raise_exc = release_raise_exc

    def __enter__(self):
        thread_id = threading.current_thread().ident
        logger.debug(f'Attempt to acquire global lock: thread {thread_id} lock {self._name}')
        acquired = self.acquire(blocking=self._blocking)
        if self._blocking and not acquired:
            logger.debug(f'Not acquired lock, but blocking=True, thread {thread_id} lock {self._name}')
            raise EnvironmentError("Lock wasn't acquired, but blocking=True")
        if not acquired:
            logger.debug(f'Not acquired the lock, thread {thread_id} lock {self._name}')
            raise AcquireFailed
        logger.debug(f'Acquire lock success, thread {thread_id} lock {self._name}')
        return self

    def __exit__(self, exc_type=None, exc_value=None, traceback=None):
        if self._release_lock_on_transaction_commit:
            transaction.on_commit(self.release)
        else:
            self.release()

    def __call__(self, func):
        @wraps(func)
        def inner(*args, **kwds):
            # 要创建一个新的锁对象
            with self.__class__(**self.kwargs_copy):
                return func(*args, **kwds)
        return inner

    def locked_by_me(self):
        if self.locked():
            if self.get_owner_id() == self.id:
                return True
        return False

    def release(self):
        try:
            super().release()
        except AcquireFailed as e:
            if self._release_raise_exc:
                raise e