添加缓存预热

backend/api/core/RedisPreheat.py

@@ -0,0 +1,133 @@
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import json
+import uuid
+
+import redis
+
+from api.core.Gzip import gzip_zip
+from core.basic.character import createCharacter
+from core.character.adventure import adv, get_adv_list
+from database.connect import RedisPool
+
+
+PreheatLoader = Callable[[], object]
+PreheatTask = tuple[str, PreheatLoader]
+PREHEAT_LOCK_KEY = 'dcalc:preheat:lock'
+PREHEAT_LOCK_EXPIRE_SECONDS = 15 * 60
+
+
+def _load_skills():
+    with open('./openapi/data/skills.json', encoding='utf-8') as file:
+        return json.load(file)
+
+
+def _iter_open_alters() -> list[str]:
+    alters: list[str] = []
+    for parent in adv:
+        for child in parent['children']:
+            alter = child.get('class')
+            if child.get('open') and alter:
+                alters.append(alter)
+    return alters
+
+
+def _build_character_preheat_tasks(equ_version: str = '0') -> list[PreheatTask]:
+    tasks: list[PreheatTask] = []
+    for alter in _iter_open_alters():
+        tasks.append(
+            (
+                f'dcalc:character:{alter}:{equ_version}',
+                lambda alter=alter: createCharacter(alter, equ_version).getInfo(),
+            )
+        )
+    return tasks
+
+
+def _write_preheated_value(redis_conn: redis.Redis, key: str, value: object, expire: int = 0, without_gzip: bool = False):
+    if value is None:
+        return
+    if without_gzip:
+        redis_conn.set(key, json.dumps(value))
+    else:
+        redis_conn.set(key, gzip_zip(json.dumps(value)))
+    if expire > 0:
+        redis_conn.expire(key, expire)
+
+
+def _run_preheat_task(key: str, loader: PreheatLoader) -> tuple[str, object]:
+    return key, loader()
+
+
+def delete_keys_with_prefix(redis_conn: redis.Redis, prefix: str):
+    keys = redis_conn.keys(f'{prefix}*')
+    if keys:
+        redis_conn.delete(*keys)
+
+
+def _acquire_preheat_lock(redis_conn: redis.Redis) -> str | None:
+    lock_token = str(uuid.uuid4())
+    locked = redis_conn.set(PREHEAT_LOCK_KEY, lock_token, nx=True, ex=PREHEAT_LOCK_EXPIRE_SECONDS)
+    if not locked:
+        return None
+    return lock_token
+
+
+def _release_preheat_lock(redis_conn: redis.Redis, lock_token: str):
+    redis_conn.eval(
+        """
+        if redis.call('GET', KEYS[1]) == ARGV[1] then
+            return redis.call('DEL', KEYS[1])
+        end
+        return 0
+        """,
+        1,
+        PREHEAT_LOCK_KEY,
+        lock_token,
+    )
+
+
+def preheat_redis_cache() -> list[str]:
+    # if config.DEBUG_MODE:
+    #     print('Skip redis preheat in debug mode.')
+    #     return []
+    base_preheat_tasks: list[PreheatTask] = [
+        ('dcalc:adventure', get_adv_list),
+        ('openapi:skills', _load_skills),
+    ]
+    character_preheat_tasks = _build_character_preheat_tasks()
+    warmed_keys: list[str] = []
+
+    with redis.Redis(connection_pool=RedisPool) as redis_conn:
+        lock_token = _acquire_preheat_lock(redis_conn)
+        if not lock_token:
+            print('Skip redis preheat: another instance is warming cache.')
+            return []
+
+        try:
+            delete_keys_with_prefix(redis_conn, 'dcalc:')
+            delete_keys_with_prefix(redis_conn, 'openapi:')
+
+            for key, loader in base_preheat_tasks:
+                try:
+                    _write_preheated_value(redis_conn, key, loader())
+                    warmed_keys.append(key)
+                except Exception as exc:
+                    print(f'Redis preheat failed for {key}: {exc}')
+
+            max_workers = min(8, len(character_preheat_tasks))
+            if max_workers > 0:
+                with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                    future_map = {executor.submit(_run_preheat_task, key, loader): key for key, loader in character_preheat_tasks}
+                    for future in as_completed(future_map):
+                        key = future_map[future]
+                        try:
+                            _, value = future.result()
+                            _write_preheated_value(redis_conn, key, value)
+                            warmed_keys.append(key)
+                        except Exception as exc:
+                            print(f'Redis preheat failed for {key}: {exc}')
+        finally:
+            _release_preheat_lock(redis_conn, lock_token)
+
+    return warmed_keys