Add distributed lock for scheduled task

Author: wu-clan

backend/app/task/service/scheduler_service.py

@@ -65,14 +65,7 @@ async def create(*, obj: CreateTaskSchedulerParam) -> None:
             if task_scheduler:
                 raise errors.ConflictError(msg='任务调度已存在')
             if obj.type == TaskSchedulerType.CRONTAB:
-                crontab_split = obj.crontab.split(' ')
-                if len(crontab_split) != 5:
-                    raise errors.RequestError(msg='Crontab 表达式非法')
-                crontab_verify('m', crontab_split[0])
-                crontab_verify('h', crontab_split[1])
-                crontab_verify('dow', crontab_split[2])
-                crontab_verify('dom', crontab_split[3])
-                crontab_verify('moy', crontab_split[4])
+                crontab_verify(obj.crontab)
             await task_scheduler_dao.create(db, obj)
 
     @staticmethod
@@ -92,14 +85,7 @@ async def update(*, pk: int, obj: UpdateTaskSchedulerParam) -> int:
                 if await task_scheduler_dao.get_by_name(db, obj.name):
                     raise errors.ConflictError(msg='任务调度已存在')
             if task_scheduler.type == TaskSchedulerType.CRONTAB:
-                crontab_split = obj.crontab.split(' ')
-                if len(crontab_split) != 5:
-                    raise errors.RequestError(msg='Crontab 表达式非法')
-                crontab_verify('m', crontab_split[0])
-                crontab_verify('h', crontab_split[1])
-                crontab_verify('dow', crontab_split[2])
-                crontab_verify('dom', crontab_split[3])
-                crontab_verify('moy', crontab_split[4])
+                crontab_verify(obj.crontab)
             count = await task_scheduler_dao.update(db, pk, obj)
             return count
 

backend/app/task/utils/schedulers.py

@@ -9,6 +9,7 @@
 
 from celery import current_app, schedules
 from celery.beat import ScheduleEntry, Scheduler
+from celery.signals import beat_init
 from celery.utils.log import get_logger
 from sqlalchemy import select
 from sqlalchemy.exc import DatabaseError, InterfaceError
@@ -28,6 +29,35 @@
 # 此计划程序必须比常规的 5 分钟更频繁地唤醒，因为它需要考虑对计划的外部更改
 DEFAULT_MAX_INTERVAL = 5  # seconds
 
+# 计划锁定时长，避免重复运行
+DEFAULT_MAX_LOCK = 300  # seconds
+
+# Copied from:
+# https://github.com/andymccurdy/redis-py/blob/master/redis/lock.py#L33
+# Changes:
+#     The second line from the bottom: The original Lua script intends
+#     to extend time to (lock remaining time + additional time); while
+#     the script here extend time to an expected expiration time.
+# KEYS[1] - lock name
+# ARGS[1] - token
+# ARGS[2] - additional milliseconds
+# return 1 if the locks time was extended, otherwise 0
+LUA_EXTEND_TO_SCRIPT = """
+    local token = redis.call('get', KEYS[1])
+    if not token or token ~= ARGV[1] then
+        return 0
+    end
+    local expiration = redis.call('pttl', KEYS[1])
+    if not expiration then
+        expiration = 0
+    end
+    if expiration < 0 then
+        return 0
+    end
+    redis.call('pexpire', KEYS[1], ARGV[2])
+    return 1
+"""
+
 logger = get_logger('fba.schedulers')
 
 
@@ -188,21 +218,12 @@ async def to_model_schedule(name: str, task: str, schedule: schedules.schedule |
                 if not obj:
                     obj = TaskScheduler(**CreateTaskSchedulerParam(task=task, **spec).model_dump())
             elif isinstance(schedule, schedules.crontab):
-                crontab_minute = schedule._orig_minute if crontab_verify('m', schedule._orig_minute, False) else '*'
-                crontab_hour = schedule._orig_hour if crontab_verify('h', schedule._orig_hour, False) else '*'
-                crontab_day_of_week = (
-                    schedule._orig_day_of_week if crontab_verify('dom', schedule._orig_day_of_week, False) else '*'
-                )
-                crontab_day_of_month = (
-                    schedule._orig_day_of_month if crontab_verify('dom', schedule._orig_day_of_month, False) else '*'
-                )
-                crontab_month_of_year = (
-                    schedule._orig_month_of_year if crontab_verify('moy', schedule._orig_month_of_year, False) else '*'
-                )
+                crontab = f'{schedule._orig_minute} {schedule._orig_hour} {schedule._orig_day_of_week} {schedule._orig_day_of_month} {schedule._orig_month_of_year}'  # noqa: E501
+                crontab_verify(crontab)
                 spec = {
                     'name': name,
                     'type': TaskSchedulerType.CRONTAB.value,
-                    'crontab': f'{crontab_minute} {crontab_hour} {crontab_day_of_week} {crontab_day_of_month} {crontab_month_of_year}',  # noqa: E501
+                    'crontab': crontab,
                 }
                 stmt = select(TaskScheduler).filter_by(**spec)
                 query = await db.execute(stmt)
@@ -269,13 +290,18 @@ def _unpack_options(
 
 
 class DatabaseScheduler(Scheduler):
+    """数据库调度程序"""
+
     Entry = ModelEntry
 
     _schedule = None
     _last_update = None
     _initial_read = True
     _heap_invalidated = False
 
+    lock = None
+    lock_key = f'{settings.CELERY_REDIS_PREFIX}:beat_lock'
+
     def __init__(self, *args, **kwargs):
         self.app = kwargs['app']
         self._dirty = set()
@@ -324,6 +350,16 @@ def reserve(self, entry):
         self._dirty.add(new_entry.name)
         return new_entry
 
+    def close(self):
+        """重写父函数"""
+        if self.lock:
+            logger.info('beat: Releasing lock')
+            if run_await(self.lock.owned)():
+                run_await(self.lock.release)()
+            self.lock = None
+
+        self.sync()
+
     def sync(self):
         """重写父函数"""
         _tried = set()
@@ -410,3 +446,29 @@ def schedule(self) -> dict[str, ModelEntry]:
 
         # logger.debug(self._schedule)
         return self._schedule
+
+
+@beat_init.connect
+def acquire_distributed_beat_lock(sender=None, *args, **kwargs):
+    """
+    尝试在启动时获取锁
+
+    :param sender: 接收方应响应的发送方
+    :return:
+    """
+    scheduler = sender.scheduler
+    if not scheduler.lock_key:
+        return
+
+    logger.debug('beat: Acquiring lock...')
+    lock = redis_client.lock(
+        scheduler.lock_key,
+        timeout=DEFAULT_MAX_LOCK,
+        sleep=scheduler.max_interval,
+    )
+    # overwrite redis-py's extend script
+    # which will add additional timeout instead of extend to a new timeout
+    lock.lua_extend = redis_client.register_script(LUA_EXTEND_TO_SCRIPT)
+    run_await(lock.acquire)()
+    logger.info('beat: Acquired lock')
+    scheduler.lock = lock

backend/app/task/utils/tzcrontab.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 from datetime import datetime
-from typing import Literal
 
 from celery import schedules
 from celery.schedules import ParseException, crontab_parser
@@ -53,34 +52,22 @@ def __reduce__(self) -> tuple[type, tuple[str, str, str, str, str], None]:
         )
 
 
-def crontab_verify(filed: Literal['m', 'h', 'dow', 'dom', 'moy'], value: str, raise_exc: bool = True) -> bool:
+def crontab_verify(crontab: str) -> None:
     """
     验证 Celery crontab 表达式
 
-    :param filed: 验证的字段
-    :param value: 验证的值
-    :param raise_exc: 是否抛出异常
+    :param crontab: 计划表达式
     :return:
     """
-    valid = True
+    crontab_split = crontab.split(' ')
+    if len(crontab_split) != 5:
+        raise errors.RequestError(msg='Crontab 表达式非法')
 
     try:
-        match filed:
-            case 'm':
-                crontab_parser(60, 0).parse(value)
-            case 'h':
-                crontab_parser(24, 0).parse(value)
-            case 'dow':
-                crontab_parser(7, 0).parse(value)
-            case 'dom':
-                crontab_parser(31, 1).parse(value)
-            case 'moy':
-                crontab_parser(12, 1).parse(value)
-            case _:
-                raise errors.ServerError(msg=f'无效字段：{filed}')
+        crontab_parser(60, 0).parse(crontab_split[0])  # minute
+        crontab_parser(24, 0).parse(crontab_split[1])  # hour
+        crontab_parser(7, 0).parse(crontab_split[2])  # day_of_week
+        crontab_parser(31, 1).parse(crontab_split[3])  # day_of_month
+        crontab_parser(12, 1).parse(crontab_split[4])  # month_of_year
     except ParseException:
-        valid = False
-        if raise_exc:
-            raise errors.RequestError(msg=f'crontab 值 {value} 非法')
-
-    return valid
+        raise errors.RequestError(msg='Crontab 表达式非法')