mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-10-30 11:54:50 +08:00
263 lines
9.6 KiB
Python
263 lines
9.6 KiB
Python
"""
|
||
内存回收装饰器模块
|
||
提供装饰器用于在函数执行后立即回收内存
|
||
"""
|
||
import ctypes
|
||
import functools
|
||
import gc
|
||
import logging
|
||
import os
|
||
import sys
|
||
from typing import Callable, Optional, ParamSpec, TypeVar, cast
|
||
|
||
import psutil
|
||
|
||
_Params = ParamSpec("_Params")
|
||
_Return = TypeVar("_Return")
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def memory_gc(
|
||
force_collect: bool = True,
|
||
log_memory_usage: bool = False,
|
||
) -> Callable[[Callable[_Params, _Return]], Callable[_Params, _Return]]:
|
||
"""
|
||
内存回收装饰器。
|
||
|
||
:param force_collect: 是否强制执行垃圾回收，默认True
|
||
:param log_memory_usage: 是否记录内存使用日志，默认False
|
||
:return: 装饰器函数
|
||
"""
|
||
def decorator(func: Callable[_Params, _Return]) -> Callable[_Params, _Return]:
|
||
"""
|
||
保留原调用参数与返回值类型，同时附加当前配置的内存处理。
|
||
"""
|
||
@functools.wraps(func)
|
||
def wrapper(*args: _Params.args, **kwargs: _Params.kwargs) -> _Return:
|
||
"""
|
||
函数退出时执行 finally 回收，原结果或异常继续交给调用方。
|
||
"""
|
||
memory_before = None
|
||
memory_after = None
|
||
if log_memory_usage:
|
||
memory_before = get_memory_usage()
|
||
logger.info(f"函数 {func.__name__} 执行前内存使用: {memory_before}")
|
||
|
||
try:
|
||
result = func(*args, **kwargs)
|
||
|
||
if log_memory_usage:
|
||
memory_after = get_memory_usage()
|
||
logger.info(f"函数 {func.__name__} 执行后内存使用: {memory_after}")
|
||
if memory_before:
|
||
memory_diff = memory_after - memory_before
|
||
logger.info(f"函数 {func.__name__} 内存变化: {memory_diff} MB")
|
||
|
||
return result
|
||
|
||
finally:
|
||
if force_collect:
|
||
collected = gc.collect()
|
||
if log_memory_usage:
|
||
logger.info(f"函数 {func.__name__} 垃圾回收完成，回收对象数: {collected}")
|
||
|
||
if log_memory_usage:
|
||
memory_after_gc = get_memory_usage()
|
||
logger.info(f"函数 {func.__name__} 垃圾回收后内存使用: {memory_after_gc}")
|
||
if memory_after:
|
||
memory_freed = memory_after - memory_after_gc
|
||
logger.info(f"函数 {func.__name__} 释放内存: {memory_freed} MB")
|
||
|
||
return wrapper
|
||
return decorator
|
||
|
||
|
||
def get_memory_usage() -> float:
|
||
"""
|
||
获取当前进程的内存使用情况（MB）。
|
||
|
||
:return: 当前进程常驻内存，单位为 MB；读取失败时返回 0.0
|
||
"""
|
||
try:
|
||
process = psutil.Process(os.getpid())
|
||
memory_info = process.memory_info()
|
||
rss: int = memory_info.rss
|
||
# 转换为MB
|
||
return rss / 1024 / 1024
|
||
except Exception as e:
|
||
logger.warning(f"获取内存使用情况失败: {e}")
|
||
return 0.0
|
||
|
||
|
||
# jemalloc 的 MALLCTL_ARENAS_ALL：对全部 arena 执行 purge
|
||
_JEMALLOC_ARENAS_ALL_PURGE = b"arena.4096.purge"
|
||
# 把调用线程的线程缓存（tcache）交还所属 arena
|
||
_JEMALLOC_THREAD_TCACHE_FLUSH = b"thread.tcache.flush"
|
||
_JEMALLOC_BACKGROUND_THREAD = b"background_thread"
|
||
_JEMALLOC_MAX_BACKGROUND_THREADS = b"max_background_threads"
|
||
# MALLOC_CONF 中与后台线程相关、不能被子进程继承的配置项
|
||
_BACKGROUND_THREAD_OPTIONS = frozenset({"background_thread", "max_background_threads"})
|
||
|
||
|
||
@functools.cache
|
||
def _process_libc() -> Optional[ctypes.CDLL]:
|
||
"""
|
||
返回当前进程的全局符号表；非 Linux 或无法加载时返回 None。
|
||
"""
|
||
if not sys.platform.startswith("linux"):
|
||
return None
|
||
try:
|
||
# CDLL(None) 取当前进程的全局符号表，包含 LD_PRELOAD 进来的 jemalloc
|
||
return ctypes.CDLL(None)
|
||
except OSError:
|
||
return None
|
||
|
||
|
||
@functools.cache
|
||
def _jemalloc_mallctl() -> Optional[Callable[..., int]]:
|
||
"""
|
||
返回 jemalloc 的 ``mallctl``；进程未加载 jemalloc 时返回 None。只解析一次，供高频调用复用。
|
||
"""
|
||
libc = _process_libc()
|
||
if libc is None:
|
||
return None
|
||
# 动态符号：仅 jemalloc 提供
|
||
mallctl = getattr(libc, "mallctl", None)
|
||
if mallctl is None:
|
||
return None
|
||
mallctl.argtypes = [ctypes.c_char_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t]
|
||
mallctl.restype = ctypes.c_int
|
||
return cast(Callable[..., int], mallctl)
|
||
|
||
|
||
def flush_thread_allocator_cache() -> bool:
|
||
"""
|
||
把调用线程在 jemalloc 线程缓存中囤积的空闲块交还 arena。
|
||
|
||
jemalloc 只在线程自身继续分配时才回收它的线程缓存，线程池里空闲等待的 worker
|
||
不再分配，最后一个任务留下的缓存会一直计入 RSS（线程越多越明显）。交还后的空闲页
|
||
再由 arena 的衰减机制归还系统。只能由持有缓存的线程自己调用；未使用 jemalloc 时不做任何事。
|
||
|
||
:return: 是否实际执行了刷新
|
||
"""
|
||
mallctl = _jemalloc_mallctl()
|
||
if mallctl is None:
|
||
return False
|
||
return mallctl(_JEMALLOC_THREAD_TCACHE_FLUSH, None, None, None, 0) == 0
|
||
|
||
|
||
def configure_allocator_background_thread(max_threads: int = 1) -> bool:
|
||
"""
|
||
让 jemalloc 后台回收线程只在主进程运行，并从子进程将继承的 ``MALLOC_CONF`` 中移除相关项。
|
||
|
||
后台线程按衰减时间把空闲脏页归还系统，但 ``MALLOC_CONF`` 与 ``LD_PRELOAD`` 会被所有子进程
|
||
继承，Chromium 的 GPU 进程在后台线程开启时无法启动，浏览器仿真整体 abort。
|
||
|
||
:param max_threads: 后台线程上限
|
||
:return: 是否在运行时实际开启了后台线程
|
||
"""
|
||
configured = _strip_background_thread_conf()
|
||
if configured:
|
||
return False
|
||
mallctl = _jemalloc_mallctl()
|
||
if mallctl is None:
|
||
return False
|
||
limit = ctypes.c_size_t(max_threads)
|
||
if mallctl(_JEMALLOC_MAX_BACKGROUND_THREADS, None, None, ctypes.byref(limit), ctypes.sizeof(limit)) != 0:
|
||
return False
|
||
enabled = ctypes.c_bool(True)
|
||
return mallctl(_JEMALLOC_BACKGROUND_THREAD, None, None, ctypes.byref(enabled), ctypes.sizeof(enabled)) == 0
|
||
|
||
|
||
def _strip_background_thread_conf() -> bool:
|
||
"""
|
||
从 ``os.environ["MALLOC_CONF"]`` 去掉后台线程项，返回原配置是否包含这些项。
|
||
"""
|
||
conf = os.environ.get("MALLOC_CONF")
|
||
if not conf:
|
||
return False
|
||
options = [option for option in conf.split(",") if option]
|
||
kept = [option for option in options if option.split(":", 1)[0] not in _BACKGROUND_THREAD_OPTIONS]
|
||
if len(kept) == len(options):
|
||
return False
|
||
if kept:
|
||
os.environ["MALLOC_CONF"] = ",".join(kept)
|
||
else:
|
||
del os.environ["MALLOC_CONF"]
|
||
return True
|
||
|
||
|
||
def release_allocator_memory() -> Optional[str]:
|
||
"""
|
||
把 C 分配器中已释放但尚未归还的空闲页交还操作系统。
|
||
|
||
``gc.collect()`` 只把 Python 对象交还分配器，jemalloc 默认没有后台回收线程，空闲 arena
|
||
的脏页会一直计入 RSS，glibc 的堆顶碎片同理。官方镜像通过 ``LD_PRELOAD`` 加载 jemalloc 时
|
||
调用 ``mallctl`` purge；源码部署等使用 glibc 的环境调用 ``malloc_trim(0)``。两者都只归还
|
||
空闲页、不影响存活对象，其他平台或符号不存在时不做任何事。
|
||
|
||
:return: 实际使用的机制（``jemalloc`` 或 ``glibc``），未执行时返回 None
|
||
"""
|
||
libc = _process_libc()
|
||
if libc is None:
|
||
return None
|
||
mallctl = _jemalloc_mallctl()
|
||
if mallctl is not None:
|
||
if mallctl(_JEMALLOC_ARENAS_ALL_PURGE, None, None, None, 0) == 0:
|
||
return "jemalloc"
|
||
return None
|
||
# 动态符号：仅 glibc 提供
|
||
malloc_trim = getattr(libc, "malloc_trim", None)
|
||
if malloc_trim is not None:
|
||
malloc_trim.argtypes = [ctypes.c_size_t]
|
||
malloc_trim.restype = ctypes.c_int
|
||
malloc_trim(0)
|
||
return "glibc"
|
||
return None
|
||
|
||
|
||
def memory_monitor(
|
||
threshold_mb: Optional[float] = None,
|
||
) -> Callable[[Callable[_Params, _Return]], Callable[_Params, _Return]]:
|
||
"""
|
||
内存监控装饰器，当内存使用超过阈值时自动触发垃圾回收。
|
||
|
||
:param threshold_mb: 内存阈值，单位为 MB；超过时触发回收，None 或 0 表示不按阈值回收
|
||
:return: 装饰器函数
|
||
"""
|
||
def decorator(func: Callable[_Params, _Return]) -> Callable[_Params, _Return]:
|
||
"""
|
||
保留原调用参数与返回值类型，同时附加当前配置的内存处理。
|
||
"""
|
||
@functools.wraps(func)
|
||
def wrapper(*args: _Params.args, **kwargs: _Params.kwargs) -> _Return:
|
||
"""
|
||
按阈值检查执行前后占用，函数抛错时保留原来的直接传播路径。
|
||
"""
|
||
current_memory = get_memory_usage()
|
||
|
||
if threshold_mb and current_memory > threshold_mb:
|
||
logger.warning(f"内存使用超过阈值 {threshold_mb}MB，当前使用: {current_memory}MB")
|
||
collected = gc.collect()
|
||
logger.info(f"自动垃圾回收完成，回收对象数: {collected}")
|
||
|
||
result = func(*args, **kwargs)
|
||
|
||
if threshold_mb:
|
||
memory_after = get_memory_usage()
|
||
if memory_after > threshold_mb:
|
||
collected = gc.collect()
|
||
logger.info(f"函数执行后垃圾回收完成，回收对象数: {collected}")
|
||
|
||
return result
|
||
|
||
return wrapper
|
||
return decorator
|
||
|
||
|
||
# 便捷的装饰器别名
|
||
memory_cleanup = memory_gc
|
||
auto_gc = memory_gc(force_collect=True, log_memory_usage=True)
|
||
memory_watch = memory_monitor
|