From 579c2c1d7a9f503dad3f674d5efc09f11b2d3407 Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Tue, 29 Sep 2020 13:33:53 +0800
Subject: [PATCH] =?UTF-8?q?feat(celery):=20=E4=BF=9D=E8=AF=81=E5=90=8C?=
 =?UTF-8?q?=E6=97=B6=E5=8F=AA=E6=9C=89=E4=B8=80=E4=B8=AAbeat=E5=9C=A8?=
 =?UTF-8?q?=E8=BF=90=E8=A1=8C=20(#4723)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat(celery): 保证同时只有一个beat在运行

* fix: 修复代码拼写错误

* fix: 修复拼写

* fix: remove import

Co-authored-by: ibuler <ibuler@qq.com>
---
 jms                        | 18 +++++-------------
 utils/start_celery_beat.py | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 13 deletions(-)
 create mode 100644 utils/start_celery_beat.py

diff --git a/jms b/jms
index 1b6935858..3969b5418 100755
--- a/jms
+++ b/jms
@@ -271,19 +271,11 @@ def get_start_flower_kwargs():
 
 def get_start_beat_kwargs():
     print("\n- Start Beat as Periodic Task Scheduler")
-    os.environ.setdefault('PYTHONOPTIMIZE', '1')
-    if os.getuid() == 0:
-        os.environ.setdefault('C_FORCE_ROOT', '1')
-
-    scheduler = "django_celery_beat.schedulers:DatabaseScheduler"
+    utils_dir = os.path.join(BASE_DIR, 'utils')
     cmd = [
-        'celery', 'beat',
-        '-A', 'ops',
-        '-l', 'INFO',
-        '--scheduler', scheduler,
-        '--max-interval', '60'
+        sys.executable, 'start_celery_beat.py',
     ]
-    return {"cmd": cmd, 'cwd': APPS_DIR}
+    return {"cmd": cmd, 'cwd': utils_dir}
 
 
 processes = {}
@@ -299,7 +291,7 @@ def watch_services():
         for s, p in processes.items():
             print("{} Check service status: {} -> ".format(now, s), end='')
             try:
-                p.wait(timeout=1)
+                p.wait(timeout=1)  # 不wait，子进程可能无法回收
             except subprocess.TimeoutExpired:
                 pass
             ok = is_running(s)
@@ -400,7 +392,7 @@ def start_service(s):
 
 def start_services_and_watch(s):
     logging.info(time.ctime())
-    logging.info('Jumpserver version {}, more see https://www.jumpserver.org'.format(
+    logging.info('JumpServer version {}, more see https://www.jumpserver.org'.format(
         __version__)
     )
 
diff --git a/utils/start_celery_beat.py b/utils/start_celery_beat.py
new file mode 100644
index 000000000..466624e3e
--- /dev/null
+++ b/utils/start_celery_beat.py
@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+#
+import os
+import sys
+import subprocess
+
+import redis_lock
+from redis import Redis
+
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+APPS_DIR = os.path.join(BASE_DIR, 'apps')
+
+sys.path.insert(0, BASE_DIR)
+from apps.jumpserver.const import CONFIG
+
+os.environ.setdefault('PYTHONOPTIMIZE', '1')
+if os.getuid() == 0:
+    os.environ.setdefault('C_FORCE_ROOT', '1')
+
+redis = Redis(host=CONFIG.REDIS_HOST, port=CONFIG.REDIS_PORT, password=CONFIG.REDIS_PASSWORD)
+scheduler = "django_celery_beat.schedulers:DatabaseScheduler"
+
+cmd = [
+    'celery', 'beat',
+    '-A', 'ops',
+    '-l', 'INFO',
+    '--scheduler', scheduler,
+    '--max-interval', '60'
+]
+
+with redis_lock.Lock(redis, name="beat-distribute-start-lock", expire=60, auto_renewal=True):
+    print("Get beat lock start to run it")
+    code = subprocess.call(cmd, cwd=APPS_DIR)
+    sys.exit(code)