59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
import uuid
|
|
from abc import ABC, abstractmethod
|
|
from threading import current_thread
|
|
|
|
from django_redis import get_redis_connection
|
|
|
|
import logging
|
|
|
|
logger = logging.getLogger('apps')
|
|
|
|
|
|
class ILock(ABC):
|
|
|
|
@abstractmethod
|
|
def try_lock(self, timeout: int) -> bool:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def unlock(self) -> None:
|
|
pass
|
|
|
|
|
|
class SimpleLock(ILock):
|
|
KEY_PREFIX = "lock:"
|
|
# ID_PREFIX = str(uuid.uuid4()).replace("-", '') + "-"
|
|
ID_PREFIX = '-'
|
|
|
|
def __init__(self, name: str, alias: str = "default") -> None:
|
|
self.name = name
|
|
self.redis_con = get_redis_connection(alias)
|
|
self.key = f"{self.KEY_PREFIX}{name}{self.ID_PREFIX}"
|
|
|
|
def try_lock(self, timeout: int) -> bool:
|
|
current_thread_id, value = self.get_value()
|
|
logger.info(f"线程:{current_thread_id}开始获取锁")
|
|
is_locked = self.set_if_absent(self.key, value, timeout)
|
|
logger.info(f"线程:{current_thread_id}获取锁是否成功: {bool(is_locked)}")
|
|
return bool(is_locked)
|
|
|
|
def get_value(self) -> tuple:
|
|
current_thread_id = current_thread().ident
|
|
value = f"{self.ID_PREFIX}{current_thread_id}"
|
|
return current_thread_id, value
|
|
|
|
def unlock(self) -> None:
|
|
current_thread_id, value = self.get_value()
|
|
logger.info(f"线程:{current_thread_id}开始删除锁")
|
|
lua_script = """if (redis.call('get', KEYS[1]) == ARGV[1]) then
|
|
-- 释放锁 del key
|
|
return redis.call('del', KEYS[1])
|
|
end
|
|
return 0
|
|
"""
|
|
deleted = self.redis_con.eval(lua_script, 1, self.key, value)
|
|
logger.info(f"线程:{current_thread_id}删除锁是否成功: {bool(deleted)}")
|
|
|
|
def set_if_absent(self, key: str, value: str, timeout: int) -> bool:
|
|
return bool(self.redis_con.set(key, value, ex=timeout, nx=True))
|