134 lines
4.2 KiB
Python
134 lines
4.2 KiB
Python
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
|