241 lines
8.4 KiB
Python
241 lines
8.4 KiB
Python
"""缓存工具 - RefreshableCache + 全局缓存注册
|
||
|
||
支持文件持久化，所有 Gunicorn worker 共享缓存。
|
||
"""
|
||
import os
|
||
import pickle
|
||
import hashlib
|
||
import tempfile
|
||
import time
|
||
from functools import wraps
|
||
|
||
_registry = []
|
||
|
||
CACHE_DIR = None
|
||
DATA_VERSION_FILE = '.data-version'
|
||
|
||
|
||
def _get_cache_dir():
|
||
"""获取文件缓存目录（项目根目录/tmp_cache）"""
|
||
global CACHE_DIR
|
||
if CACHE_DIR is not None:
|
||
return CACHE_DIR
|
||
# 本项目在 services/ 下，项目根是 services/..
|
||
base = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
CACHE_DIR = os.path.join(base, 'tmp_cache')
|
||
os.makedirs(CACHE_DIR, exist_ok=True)
|
||
return CACHE_DIR
|
||
|
||
|
||
def _data_version_signature():
|
||
"""返回业务数据版本文件签名，供所有进程识别导入后的数据变化。"""
|
||
path = os.path.join(_get_cache_dir(), DATA_VERSION_FILE)
|
||
try:
|
||
stat = os.stat(path)
|
||
# 同一纳秒时间戳下连续写入仍可能发生，加入内容摘要避免漏判。
|
||
with open(path, 'rb') as f:
|
||
content = f.read()
|
||
return (stat.st_ino, stat.st_mtime_ns, stat.st_size,
|
||
hashlib.sha256(content).digest())
|
||
except OSError:
|
||
return None
|
||
|
||
|
||
def bump_data_version():
|
||
"""发布一次业务数据变更，使所有 worker 的旧缓存立即失效。"""
|
||
cache_dir = _get_cache_dir()
|
||
path = os.path.join(cache_dir, DATA_VERSION_FILE)
|
||
fd, tmp_path = tempfile.mkstemp(dir=cache_dir, suffix='.version.tmp')
|
||
try:
|
||
with os.fdopen(fd, 'w', encoding='utf-8') as f:
|
||
f.write(str(time.time_ns()))
|
||
os.replace(tmp_path, path)
|
||
except Exception:
|
||
try:
|
||
os.unlink(tmp_path)
|
||
except OSError:
|
||
pass
|
||
raise
|
||
return _data_version_signature()
|
||
|
||
|
||
class RefreshableCache:
|
||
"""可刷新的缓存（内存 + 文件持久化）
|
||
|
||
所有 Gunicorn worker 共享文件缓存。
|
||
数据写入 auto_update.py → 删除文件缓存 → 各 worker 下次请求自动重建。
|
||
"""
|
||
|
||
def __init__(self, maxsize=1, cache_key=None):
|
||
self.maxsize = maxsize
|
||
self._cache = None
|
||
self._file_signature = None
|
||
self._data_signature = None
|
||
self._cache_key = cache_key
|
||
self._file_path = None
|
||
self._wrapped_func = None
|
||
|
||
def _current_file_signature(self):
|
||
"""返回共享缓存文件签名，用于识别其他进程的刷新/重建。"""
|
||
fpath = self._file_path
|
||
if not fpath:
|
||
return None
|
||
try:
|
||
stat = os.stat(fpath)
|
||
# os.replace 会生成新 inode；即使新旧内容大小相同、写入时间极近也能识别。
|
||
return (stat.st_ino, stat.st_mtime_ns, stat.st_size)
|
||
except OSError:
|
||
return None
|
||
|
||
def _get_file_path(self, func):
|
||
"""获取缓存文件路径（懒计算）"""
|
||
if self._file_path:
|
||
return self._file_path
|
||
key = self._cache_key
|
||
if not key:
|
||
key = f"{func.__module__}.{func.__qualname__}"
|
||
fname = hashlib.md5(key.encode()).hexdigest() + '.cache'
|
||
self._file_path = os.path.join(_get_cache_dir(), fname)
|
||
return self._file_path
|
||
|
||
def _load_from_file(self):
|
||
"""从文件加载缓存；数据版本晚于缓存文件时拒绝旧结果。"""
|
||
fpath = self._file_path
|
||
if not fpath or not os.path.exists(fpath):
|
||
return False
|
||
try:
|
||
version_path = os.path.join(_get_cache_dir(), DATA_VERSION_FILE)
|
||
if os.path.exists(version_path):
|
||
cache_mtime = os.stat(fpath).st_mtime_ns
|
||
version_mtime = os.stat(version_path).st_mtime_ns
|
||
# 版本文件与缓存文件可能落在同一纳秒时间戳；版本不早于缓存时，
|
||
# 按旧缓存处理，避免导入后重新加载更新前结果。
|
||
if version_mtime >= cache_mtime:
|
||
try:
|
||
os.remove(fpath)
|
||
except OSError:
|
||
pass
|
||
return False
|
||
with open(fpath, 'rb') as f:
|
||
self._cache = pickle.load(f)
|
||
self._file_signature = self._current_file_signature()
|
||
self._data_signature = _data_version_signature()
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
def _save_to_file(self, data):
|
||
"""保存缓存到文件（原子写入）"""
|
||
fpath = self._file_path
|
||
if not fpath:
|
||
return
|
||
try:
|
||
cache_dir = os.path.dirname(fpath)
|
||
os.makedirs(cache_dir, exist_ok=True)
|
||
fd, tmp_path = tempfile.mkstemp(dir=cache_dir, suffix='.tmp')
|
||
try:
|
||
with os.fdopen(fd, 'wb') as f:
|
||
pickle.dump(data, f)
|
||
os.replace(tmp_path, fpath)
|
||
self._file_signature = self._current_file_signature()
|
||
self._data_signature = _data_version_signature()
|
||
except:
|
||
try:
|
||
os.unlink(tmp_path)
|
||
except OSError:
|
||
pass
|
||
except Exception:
|
||
pass
|
||
|
||
def wrap(self, func):
|
||
"""装饰器：缓存函数返回值，支持内存+文件双缓存"""
|
||
self._wrapped_func = func
|
||
self._get_file_path(func)
|
||
|
||
@wraps(func)
|
||
def wrapper(*args, **kwargs):
|
||
# 1. 内存缓存命中前核对共享文件。数据更新脚本运行在独立进程中，
|
||
# 删除/重建缓存文件后，各 Gunicorn worker 必须放弃自己的旧内存副本。
|
||
if self._cache is not None:
|
||
current_signature = self._current_file_signature()
|
||
current_data_signature = _data_version_signature()
|
||
if (current_signature == self._file_signature and current_signature is not None
|
||
and current_data_signature == self._data_signature):
|
||
return self._cache
|
||
self._cache = None
|
||
self._file_signature = None
|
||
self._data_signature = None
|
||
|
||
# 2. 文件缓存命中 → 加载到内存
|
||
if self._load_from_file():
|
||
return self._cache
|
||
|
||
# 3. 执行函数
|
||
result = func(*args, **kwargs)
|
||
self._cache = result
|
||
|
||
# 4. 写入文件缓存（异步感觉不到延迟）
|
||
self._save_to_file(result)
|
||
|
||
return result
|
||
return wrapper
|
||
|
||
def refresh(self):
|
||
"""清除缓存（内存 + 文件）"""
|
||
self._cache = None
|
||
self._file_signature = None
|
||
self._data_signature = None
|
||
fpath = self._file_path
|
||
if fpath and os.path.exists(fpath):
|
||
try:
|
||
os.remove(fpath)
|
||
except OSError:
|
||
pass
|
||
# 尝试用 cache_key 匹配更多文件（当 key 在 __init__.py 中有变化时兜底）
|
||
# 但如果 _file_path 从未初始化过则跳过
|
||
if not fpath and self._cache_key:
|
||
# 尝试从 cache_key 推导文件路径
|
||
fname = hashlib.md5(self._cache_key.encode()).hexdigest() + '.cache'
|
||
fallback = os.path.join(_get_cache_dir(), fname)
|
||
if os.path.exists(fallback):
|
||
try:
|
||
os.remove(fallback)
|
||
except OSError:
|
||
pass
|
||
|
||
def clear(self):
|
||
"""清除缓存"""
|
||
self.refresh()
|
||
|
||
|
||
def register_cache(cache):
|
||
"""注册缓存实例"""
|
||
_registry.append(cache)
|
||
|
||
|
||
def refresh_all_caches():
|
||
"""刷新所有注册的缓存（内存 + 文件），返回已刷新的缓存名称。"""
|
||
refreshed = []
|
||
for c in _registry:
|
||
c.refresh()
|
||
refreshed.append(c._cache_key or (
|
||
f"{c._wrapped_func.__module__}.{c._wrapped_func.__qualname__}"
|
||
if c._wrapped_func else c.__class__.__name__))
|
||
# 兜底：确保文件缓存目录中所有 .cache 文件被清除
|
||
cache_dir = _get_cache_dir()
|
||
try:
|
||
for fname in os.listdir(cache_dir):
|
||
if fname.endswith('.cache'):
|
||
os.remove(os.path.join(cache_dir, fname))
|
||
except Exception:
|
||
pass
|
||
# 发布数据版本。即使旧缓存文件因权限或时序问题未删除，其他 worker
|
||
# 也会在下一次请求时发现版本变化并重新计算。
|
||
bump_data_version()
|
||
return refreshed
|
||
|
||
|
||
def clear_all_caches():
|
||
"""清除所有注册的缓存"""
|
||
refresh_all_caches()
|