|
|
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
|
