From c5c043052c2ea8f628cb4363edd3bc70e95ffafc Mon Sep 17 00:00:00 2001 From: Wu Clan Date: Wed, 16 Sep 2026 16:28:35 +0800 Subject: [PATCH] Optimize Redis connection pool and batch operations (#1252) --- backend/common/cache/pubsub.py | 35 ++++++++----- backend/core/conf.py | 3 ++ backend/database/redis.py | 94 ++++++++++++++++++++++++++++++---- 3 files changed, 109 insertions(+), 23 deletions(-) diff --git a/backend/common/cache/pubsub.py b/backend/common/cache/pubsub.py index 797c6b3a..7c32ab97 100644 --- a/backend/common/cache/pubsub.py +++ b/backend/common/cache/pubsub.py @@ -38,26 +38,33 @@ class CachePubSubManager: try: # 使用独立连接 - pubsub_client = RedisCli(socket_timeout=None) + pubsub_client = RedisCli(max_connections=1) pubsub = pubsub_client.pubsub() await pubsub.subscribe(settings.CACHE_PUBSUB_CHANNEL) # 发布订阅成功 reconnect_attempts = 0 - async for message in pubsub.listen(): - if message['type'] == 'message': - try: - data = json.loads(message['data']) - cache_key = data['cache_key'] - if not data['delete_by_prefix']: - local_cache_manager.delete(cache_key) - else: - local_cache_manager.delete_by_prefix(cache_key) - except json.JSONDecodeError as e: - log.warning(f'[CachePubSub] 消息格式错误 {e}') - except Exception as e: - log.error(f'[CachePubSub] 处理通知失败: {e}') + # 带超时轮询而不是 listen()，每次进入读取都会触发健康检查 PING， + # 避免连接被静默断开后协程永久挂起 + while True: + message = await pubsub.get_message( + ignore_subscribe_messages=True, + timeout=settings.CACHE_PUBSUB_POLL_TIMEOUT, + ) + if message is None or message['type'] != 'message': + continue + try: + data = json.loads(message['data']) + cache_key = data['cache_key'] + if not data['delete_by_prefix']: + local_cache_manager.delete(cache_key) + else: + local_cache_manager.delete_by_prefix(cache_key) + except json.JSONDecodeError as e: + log.warning(f'[CachePubSub] 消息格式错误 {e}') + except Exception as e: + log.error(f'[CachePubSub] 处理通知失败: {e}') except asyncio.CancelledError: break diff --git a/backend/core/conf.py b/backend/core/conf.py index 11d09f27..3e5e87d2 100644 --- a/backend/core/conf.py +++ b/backend/core/conf.py @@ -68,6 +68,8 @@ class Settings(BaseSettings): # Redis REDIS_TIMEOUT: int = 5 + REDIS_MAX_CONNECTIONS: int = 100 # 连接池上限 + REDIS_POOL_TIMEOUT: int = 20 # 等待空闲连接超时（秒） # 缓存 CACHE_LOCAL_ENABLED: bool = True @@ -79,6 +81,7 @@ class Settings(BaseSettings): CACHE_PUBSUB_CHANNEL: str = 'fba:cache:invalidate' CACHE_PUBSUB_RECONNECT_DELAY: int = 5 # 重连延迟（秒） CACHE_PUBSUB_MAX_RECONNECT_ATTEMPTS: int = 10 # 最大重连次数 + CACHE_PUBSUB_POLL_TIMEOUT: float = 1.0 # 订阅消息轮询超时（秒） # .env Snowflake SNOWFLAKE_ENABLED: bool = False diff --git a/backend/database/redis.py b/backend/database/redis.py index 3e26af82..b8eacafc 100644 --- a/backend/database/redis.py +++ b/backend/database/redis.py @@ -1,6 +1,6 @@ import sys -from redis.asyncio import Redis +from redis.asyncio import BlockingConnectionPool, Redis from redis.exceptions import AuthenticationError, TimeoutError from backend.common.log import log @@ -22,6 +22,8 @@ class RedisCli(Redis): socket_keepalive: bool = True, health_check_interval: int = 30, decode_responses: bool = True, + max_connections: int = settings.REDIS_MAX_CONNECTIONS, + pool_timeout: int = settings.REDIS_POOL_TIMEOUT, ) -> None: """ 初始化 Redis 客户端 @@ -35,8 +37,12 @@ class RedisCli(Redis): :param socket_keepalive: 是否开启 TCP Keepalive 探测 :param health_check_interval: 健康检查间隔时间（秒） :param decode_responses: 是否自动将 Redis 返回的字节流（bytes）解码为字符串（utf-8） + :param max_connections: 连接池最大连接数，超出后排队等待而不是无限新建 + :param pool_timeout: 等待空闲连接的超时时间（秒） """ - super().__init__( + pool = BlockingConnectionPool( + max_connections=max_connections, + timeout=pool_timeout, host=host, port=port, password=password, @@ -47,6 +53,9 @@ class RedisCli(Redis): health_check_interval=health_check_interval, decode_responses=decode_responses, ) + super().__init__(connection_pool=pool) + # 连接池由客户端独占，aclose 时一并释放 + self.auto_close_connection_pool = True async def init(self) -> None: """初始化 Redis 服务器""" @@ -67,13 +76,15 @@ class RedisCli(Redis): key_prefix: str, exclude_keys: str | list[str] | None = None, batch_size: int = 1000, + count: int = 1000, ) -> None: """ 删除指定前缀的所有 key :param key_prefix: 要删除的键前缀 :param exclude_keys: 要排除的键或键列表 - :param batch_size: 批量删除的大小，避免一次性删除过多键导致 Redis 阻塞 + :param batch_size: 批量删除的大小 + :param count: 每次扫描批次的数量 :return: """ exclude_set = ( @@ -84,22 +95,18 @@ class RedisCli(Redis): else set() ) batch_keys = [] - if key_prefix not in exclude_set and await self.exists(key_prefix): batch_keys.append(key_prefix) - - async for key in self.scan_iter(match=f'{key_prefix}:*'): + async for key in self.scan_iter(match=f'{key_prefix}:*', count=count): if key not in exclude_set: batch_keys.append(key) - if len(batch_keys) >= batch_size: await self.delete(*batch_keys) batch_keys.clear() - if batch_keys: await self.delete(*batch_keys) - async def get_by_prefix(self, key_prefix: str, count: int = 100) -> list[str]: + async def get_by_prefix(self, key_prefix: str, count: int = 1000) -> list[str]: """ 获取指定前缀的所有 key @@ -109,6 +116,75 @@ class RedisCli(Redis): """ return [key async for key in self.scan_iter(match=f'{key_prefix}:*', count=count)] + async def mget_batched(self, keys: list[str], batch_size: int = 1000) -> list[str | None]: + """ + 分批获取多个 key 的值 + + :param keys: 键列表 + :param batch_size: 每批数量 + :return: + """ + if batch_size <= 0: + raise ValueError('batch_size 必须大于 0') + if not keys: + return [] + values: list[str | None] = [] + for index in range(0, len(keys), batch_size): + values.extend(await self.mget(keys[index : index + batch_size])) + return values + + async def exists_batched(self, keys: list[str], batch_size: int = 1000) -> list[bool]: + """ + 分批判断多个 key 是否存在 + + :param keys: 键列表 + :param batch_size: 每批数量 + :return: + """ + if batch_size <= 0: + raise ValueError('batch_size 必须大于 0') + return [value is not None for value in await self.mget_batched(keys, batch_size=batch_size)] + + async def smembers_many(self, keys: list[str], batch_size: int = 100) -> list[set[str]]: + """ + 分批获取多个集合的成员 + + :param keys: 键列表 + :param batch_size: 每批数量 + :return: + """ + if batch_size <= 0: + raise ValueError('batch_size 必须大于 0') + if not keys: + return [] + members: list[set[str]] = [] + for index in range(0, len(keys), batch_size): + batch = keys[index : index + batch_size] + async with self.pipeline(transaction=False) as pipe: + for key in batch: + pipe.smembers(key) + results = await pipe.execute() + members.extend(set(result) if result else set() for result in results) + return members + + async def delete_batched(self, keys: list[str], batch_size: int = 1000) -> int: + """ + 分批删除多个 key + + :param keys: 键列表 + :param batch_size: 每批数量 + :return: + """ + if batch_size <= 0: + raise ValueError('batch_size 必须大于 0') + if not keys: + return 0 + deleted = 0 + for index in range(0, len(keys), batch_size): + batch = keys[index : index + batch_size] + deleted += await self.delete(*batch) + return deleted + # 创建 redis 客户端单例 redis_client: RedisCli = RedisCli()