jumpserver/apps/orgs/lock.py

from uuid import uuid4
from functools import wraps

from django.core.cache import cache
from django.db.transaction import atomic
from rest_framework.request import Request
from rest_framework.exceptions import NotAuthenticated

from orgs.utils import current_org
from common.exceptions import SomeoneIsDoingThis, Timeout
from common.utils.timezone import dt_formater, now

# Redis 中锁值得模板，该模板提供了很强的可读性，方便调试与排错
VALUE_TEMPLATE = '{stage}:{username}:{user_id}:{now}:{rand_str}'

# 锁的状态
DOING = 'doing'  # 处理中，此状态的锁可以被干掉
COMMITING = 'commiting'  # 提交事务中，此状态很重要，要确保事务在锁消失之前返回了，不要轻易删除该锁

client = cache.client.get_client(write=True)


"""
将锁的状态从 `doing` 切换到 `commiting`
KEYS[1]： key
ARGV[1]: doingvalue
ARGV[2]: commitingvalue
ARGV[3]: timeout
"""
change_lock_state_to_commiting_lua = '''
if (redis.call("get", KEYS[1]) == ARGV[1])
then
    return redis.call("set", KEYS[1], ARGV[2], "EX", ARGV[3], "XX")
else
    return 0
end
'''
change_lock_state_to_commiting_lua_obj = client.register_script(change_lock_state_to_commiting_lua)


"""
释放锁，两种`value`都要检查`doing`和`commiting`
KEYS[1]： key
ARGV[1]: 两个 `value` 中的其中一个
ARGV[2]: 两个 `value` 中的其中一个
"""
release_lua = '''
if (redis.call("get",KEYS[1]) == ARGV[1] or redis.call("get",KEYS[1]) == ARGV[2])
then
    return redis.call("del",KEYS[1])
else
    return 0
end
'''
release_lua_obj = client.register_script(release_lua)


def acquire(key, value, timeout):
    return client.set(key, value, ex=timeout, nx=True)


def get(key):
    return client.get(key)


def change_lock_state_to_commiting(key, doingvalue, commitingvalue, timeout=600):
    # 将锁的状态从 `doing` 切换到 `commiting`
    return bool(change_lock_state_to_commiting_lua_obj(keys=(key,), args=(doingvalue, commitingvalue, timeout)))


def release(key, value1, value2):
    # 释放锁，两种`value` `doing`和`commiting` 都要检查
    return release_lua_obj(keys=(key,), args=(value1, value2))


def _generate_value(request: Request, stage=DOING):
    # 不支持匿名用户
    user = request.user
    if user.is_anonymous:
        raise NotAuthenticated

    return VALUE_TEMPLATE.format(
        stage=stage, username=user.username, user_id=user.id,
        now=dt_formater(now()), rand_str=uuid4()
    )


default_wait_msg = SomeoneIsDoingThis.default_detail


def org_level_transaction_lock(key, timeout=300, wait_msg=default_wait_msg):
    """
    被装饰的 `View` 必须取消自身的 `ATOMIC_REQUESTS`，因为该装饰器要有事务的完全控制权
    [官网](https://docs.djangoproject.com/en/3.1/topics/db/transactions/#tying-transactions-to-http-requests)

    1. 获取锁：只有当锁对应的 `key` 不存在时成功获取，`value` 设置为 `doing`
    2. 开启事务：本次请求的事务必须确保在这里开启
    3. 执行 `View` 体
    4. `View` 体执行结束未异常，此时事务还未提交
    5. 检查锁是否过时，过时事务回滚，不过时，重新设置`key`延长`key`有效期，已确保足够时间提交事务，同时把`key`的状态改为`commiting`
    6. 提交事务
    7. 释放锁，释放的时候会检查`doing`与`commiting`的值，因为删除或者更改锁必须提供与当前锁的`value`相同的值，确保不误删
       [锁参考文章](http://doc.redisfans.com/string/set.html#id2)
    """

    def decorator(fun):
        @wraps(fun)
        def wrapper(request, *args, **kwargs):
            # `key`可能是组织相关的，如果是把组织`id`加上
            _key = key.format(org_id=current_org.id)
            doing_value = _generate_value(request)
            commiting_value = _generate_value(request, stage=COMMITING)
            try:
                lock = acquire(_key, doing_value, timeout)
                if not lock:
                    raise SomeoneIsDoingThis(detail=wait_msg)
                with atomic(savepoint=False):
                    ret = fun(request, *args, **kwargs)
                    # 提交事务前，检查一下锁是否还在
                    # 锁在的话，更新锁的状态为 `commiting`，延长锁时间，确保事务提交
                    # 锁不在的话回滚
                    ok = change_lock_state_to_commiting(_key, doing_value, commiting_value)
                    if not ok:
                        # 超时或者被中断了
                        raise Timeout
                    return ret
            finally:
                # 释放锁，锁的两个值都要尝试，不确定异常是从什么位置抛出的
                release(_key, commiting_value, doing_value)
        return wrapper
    return decorator