refactor: replace try_lock and un_lock with RedisLock for improved locking mechanism

Author: liuruibin

apps/chat/serializers/chat_record.py

@@ -19,7 +19,7 @@
     ApplicationChatRecordQuerySerializers
 from common.db.search import page_search
 from common.exception.app_exception import AppApiException
-from common.utils.lock import try_lock, un_lock
+from common.utils.lock import RedisLock
 
 
 class VoteRequest(serializers.Serializer):
@@ -48,7 +48,8 @@ def vote(self, instance: Dict, with_valid=True):
         if with_valid:
             self.is_valid(raise_exception=True)
             VoteRequest(data=instance).is_valid(raise_exception=True)
-        if not try_lock(self.data.get('chat_record_id')):
+        rlock = RedisLock()
+        if not rlock.try_lock(self.data.get('chat_record_id')):
             raise AppApiException(500,
                                   gettext(
                                       "Voting on the current session minutes, please do not send repeated requests"))
@@ -75,7 +76,7 @@ def vote(self, instance: Dict, with_valid=True):
                 else:
                     raise AppApiException(500, gettext("Already voted, please cancel first and then vote again"))
         finally:
-            un_lock(self.data.get('chat_record_id'))
+            rlock.un_lock(self.data.get('chat_record_id'))
         ChatCountSerializer(data={'chat_id': self.data.get('chat_id')}).update_chat()
         return True
 

apps/common/event/__init__.py

@@ -12,6 +12,7 @@
 from .listener_manage import *
 from ..constants.cache_version import Cache_Version
 from ..db.sql_execute import update_execute
+from ..utils.lock import RedisLock
 
 update_document_status_sql = """
                              UPDATE "public"."document"
@@ -22,8 +23,8 @@
 
 def run():
     from models_provider.models import Model, Status
-
-    if try_lock('event_init', 30 * 30):
+    rlock = RedisLock()
+    if rlock.try_lock('event_init', 30 * 30):
         try:
             # 修改Model状态为ERROR
             QuerySet(Model).filter(
@@ -36,4 +37,4 @@ def run():
             version, get_key = Cache_Version.SYSTEM.value
             cache.delete(get_key(key='rsa_key'), version=version)
         finally:
-            un_lock('event_init')
+            rlock.un_lock('event_init')

apps/common/event/listener_manage.py

@@ -20,7 +20,7 @@
 from common.config.embedding_config import VectorStore
 from common.db.search import native_search, get_dynamics_model, native_update
 from common.utils.common import get_file_content
-from common.utils.lock import try_lock, un_lock
+from common.utils.lock import RedisLock
 from common.utils.logger import maxkb_logger
 from common.utils.page_utils import page_desc
 from knowledge.models import Paragraph, Status, Document, ProblemParagraphMapping, TaskType, State,SourceType, SearchMode
@@ -253,7 +253,8 @@ def embedding_by_document(document_id, embedding_model: Embeddings, state_list=N
         """
         if state_list is None:
             state_list = [State.PENDING, State.SUCCESS, State.FAILURE, State.REVOKE, State.REVOKED]
-        if not try_lock('embedding:' + str(document_id)):
+        rlock = RedisLock()
+        if not rlock.try_lock('embedding:' + str(document_id)):
             return
         try:
             def is_the_task_interrupted():
@@ -290,7 +291,7 @@ def is_the_task_interrupted():
             ListenerManagement.post_update_document_status(document_id, TaskType.EMBEDDING)
             ListenerManagement.get_aggregation_document_status(document_id)()
             maxkb_logger.info(_('End--->Embedding document: {document_id}').format(document_id=document_id))
-            un_lock('embedding:' + str(document_id))
+            rlock.un_lock('embedding:' + str(document_id))
 
     @staticmethod
     def embedding_by_knowledge(knowledge_id, embedding_model: Embeddings):

apps/common/job/clean_chat_job.py

@@ -8,7 +8,7 @@
 
 from application.models import Application, Chat, ChatRecord
 from common.job.scheduler import scheduler
-from common.utils.lock import try_lock, un_lock, lock
+from common.utils.lock import lock, RedisLock
 from common.utils.logger import maxkb_logger
 from knowledge.models import File
 
@@ -70,7 +70,8 @@ def clean_chat_log_job_lock():
 
 
 def run():
-    if try_lock('clean_chat_log_job', 30 * 30):
+    rlock = RedisLock()
+    if rlock.try_lock('clean_chat_log_job', 30 * 30):
         try:
             maxkb_logger.debug('get lock clean_chat_log_job')
 
@@ -79,4 +80,4 @@ def run():
                 existing_job.remove()
             scheduler.add_job(clean_chat_log_job, 'cron', hour='0', minute='5', id='clean_chat_log')
         finally:
-            un_lock('clean_chat_log_job')
+            rlock.un_lock('clean_chat_log_job')

apps/common/job/clean_debug_file_job.py

@@ -5,7 +5,7 @@
 from django.utils import timezone
 
 from common.job.scheduler import scheduler
-from common.utils.lock import un_lock, try_lock, lock
+from common.utils.lock import lock, RedisLock
 from common.utils.logger import maxkb_logger
 from knowledge.models import File, FileSourceType
 
@@ -25,12 +25,14 @@ def clean_debug_file_lock():
     File.objects.filter(
         Q(create_time__lt=one_days_ago, source_type=FileSourceType.TEMPORARY_1_DAY.value) |
         Q(create_time__lt=two_hours_ago, source_type=FileSourceType.TEMPORARY_120_MINUTE.value) |
-        Q(create_time__lt=minutes_30_ago, source_type=FileSourceType.TEMPORARY_30_MINUTE.value)).delete()
+        Q(create_time__lt=minutes_30_ago, source_type=FileSourceType.TEMPORARY_30_MINUTE.value)
+    ).delete()
     maxkb_logger.debug(_('end clean debug file'))
 
 
 def run():
-    if try_lock('clean_debug_file', 30 * 30):
+    rlock = RedisLock()
+    if rlock.try_lock('clean_debug_file', 30 * 30):
         try:
             maxkb_logger.debug('get lock clean_debug_file')
 
@@ -39,4 +41,4 @@ def run():
                 clean_debug_file_job.remove()
             scheduler.add_job(clean_debug_file, 'cron', hour='*', minute='*/30', second='0', id='clean_debug_file')
         finally:
-            un_lock('clean_debug_file')
+            rlock.un_lock('clean_debug_file')

apps/common/job/client_access_num_job.py

@@ -11,7 +11,7 @@
 
 from application.models import ApplicationChatUserStats
 from common.job.scheduler import scheduler
-from common.utils.lock import try_lock, un_lock, lock
+from common.utils.lock import lock, RedisLock
 from common.utils.logger import maxkb_logger
 
 
@@ -28,7 +28,8 @@ def client_access_num_reset_job_lock():
 
 
 def run():
-    if try_lock('access_num_reset', 30 * 30):
+    rlock = RedisLock()
+    if rlock.try_lock('access_num_reset', 30 * 30):
         try:
             maxkb_logger.debug('get lock access_num_reset')
 
@@ -38,4 +39,4 @@ def run():
             scheduler.add_job(client_access_num_reset_job, 'cron', hour='0', minute='0', second='0',
                               id='access_num_reset')
         finally:
-            un_lock('access_num_reset')
+            rlock.un_lock('access_num_reset')

apps/common/management/commands/services/services/gunicorn.py

@@ -1,3 +1,5 @@
+import subprocess
+
 from .base import BaseService
 from ..hands import *
 
@@ -35,3 +37,15 @@ def cmd(self):
     @property
     def cwd(self):
         return APPS_DIR
+
+    def open_subprocess(self):
+        # 复制当前环境变量，并设置 ENABLE_SCHEDULER=1
+        env = os.environ.copy()
+        env['ENABLE_SCHEDULER'] = '1'
+        kwargs = {
+            'cwd': self.cwd,
+            'stderr': self.log_file,
+            'stdout': self.log_file,
+            'env': env
+        }
+        self._process = subprocess.Popen(self.cmd, **kwargs)

apps/common/utils/lock.py

@@ -6,32 +6,47 @@
     @date：2023/9/11 11:45
     @desc:
 """
-from datetime import timedelta
+from functools import wraps
 
+import uuid_utils.compat as uuid
 from django.core.cache import caches
+from django_redis import get_redis_connection
 
 memory_cache = caches['default']
 
+class RedisLock():
+    def __init__(self):
+        self.lock_value = None
 
-def try_lock(key: str, timeout=None):
-    """
-    获取锁
-    :param key:    获取锁 key
-    :param timeout 超时时间
-    :return: 是否获取到锁
-    """
-    if timeout is None:
-        timeout = 3600  # 默认超时时间为3600秒
-    return memory_cache.add(key, 'lock', timeout=timeout)
+    def try_lock(self, key: str, timeout=None):
+        """
+        获取锁
+        :param key:    获取锁 key
+        :param timeout 超时时间
+        :return: 是否获取到锁
+        """
+        redis_client = get_redis_connection("default")
+        if timeout is None:
+            timeout = 3600  # 默认超时时间为3600秒
+        self.lock_value = str(uuid.uuid7())
+        return redis_client.set(key, self.lock_value, nx=True, ex=timeout)
 
 
-def un_lock(key: str):
-    """
-    解锁
-    :param key: 解锁 key
-    :return: 是否解锁成功
-    """
-    return memory_cache.delete(key)
+    def un_lock(self, key: str):
+        """
+        解锁
+        :param key: 解锁 key
+        :return: 是否解锁成功
+        """
+        redis_client = get_redis_connection("default")
+        unlock_script = """
+            if redis.call("get", KEYS[1]) == ARGV[1] then
+                return redis.call("del", KEYS[1])
+            else
+                return 0
+            end
+            """
+        redis_client.eval(unlock_script, 1, key, self.lock_value)
 
 
 def lock(lock_key, timeout=None):
@@ -43,15 +58,19 @@ def lock(lock_key, timeout=None):
 
     """
 
-    def inner(func):
-        def run(*args, **kwargs):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
             key = lock_key(*args, **kwargs) if callable(lock_key) else lock_key
+            rlock = RedisLock()
+            if not rlock.try_lock(key, timeout):
+                # 获取锁失败，可自定义异常或返回
+                return None
             try:
-                if try_lock(key=key, timeout=timeout):
-                    return func(*args, **kwargs)
+                return func(*args, **kwargs)
             finally:
-                un_lock(key=key)
+                rlock.un_lock(key)
 
-        return run
+        return wrapper
 
-    return inner
+    return decorator

apps/maxkb/wsgi.py

@@ -26,4 +26,8 @@ def post_handler():
     job.run()
     DatabaseModelManage.init()
 
-post_handler()
+# 仅在web中启动定时任务，local_model celery 不需要
+if os.environ.get('ENABLE_SCHEDULER') == '1':
+    post_handler()
+
+