mirror of https://github.com/jumpserver/jumpserver
132 lines
4.7 KiB
Python
132 lines
4.7 KiB
Python
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
|