25.09.10

favicon_app/utils/redis_pool.py

@@ -29,34 +29,43 @@ async def get_redis() -> AsyncGenerator[Redis, None]:
         yield conn
 
 
-async def set_cache(key: str, value: [str | int], ttl: int = None) -> None:
+async def set_cache(key: str, value: [str | int], ttl: int = None, prefix: str = None) -> None:
     if not key:
         return
 
     try:
         async for redis in get_redis():
-            await redis.set(key, value, ex=ttl)
+            _key = key
+            if prefix:
+                _key = f"{prefix}{key}"
+                await redis.sadd(prefix, key)
+                await redis.expire(prefix, ttl)
+            await redis.set(_key, value, ex=ttl)
     except Exception as e:
         logger.error(f"存入redis时出错：{e}")
 
 
-async def get_cache(key: str) -> Optional[str | int]:
+async def get_cache(key: str, prefix: str = None) -> Optional[str | int]:
     if not key:
         return None
 
     try:
         async for redis in get_redis():
+            if prefix:
+                key = f"{prefix}{key}"
             return await redis.get(key)
     except Exception as e:
         logger.error(f"读取redis时出错：{e}")
 
 
-async def exist_cache(key: str) -> bool:
+async def exist_cache(key: str, prefix: str = None) -> bool:
     if not key:
         return False
 
     try:
         async for redis in get_redis():
+            if prefix:
+                key = f"{prefix}{key}"
             result = await redis.exists(key)
             return result > 0
     except Exception as e:
@@ -64,80 +73,62 @@ async def exist_cache(key: str) -> bool:
         return False
 
 
-async def remove_cache(key: str) -> None:
+async def remove_cache(key: str, prefix: str = None) -> None:
     if not key:
         return
 
     try:
         async for redis in get_redis():
-            await redis.delete(key)
+            _key = key
+            if prefix:
+                _key = f"{prefix}{key}"
+                await redis.srem(prefix, key)
+            await redis.delete(_key)
     except Exception as e:
         logger.error(f"删除redis时出错：{e}")
 
 
-async def get_cache_size(cache_name: str = "default") -> int:
+async def get_cache_size(prefix: str = None) -> int:
+    """根据前缀统计数量，用于统计Set集合
+    """
     try:
         async for redis in get_redis():
-            return await redis.llen(cache_name)
+            return await redis.scard(prefix)
     except Exception as e:
         logger.error(f"获取队列大小时出错：{e}")
         return 0
 
 
-async def set_failed_domain(domain: str, expire_seconds: int = setting.time_of_7_days) -> None:
-    """将失败的域名存入Redis，并设置过期时间
-
-    Args:
-        domain: 失败的域名
-        expire_seconds: 过期时间（秒），默认为7天
-    """
+async def set_failed_domain(domain: str, expire_seconds: int = None) -> None:
     if not domain:
         return
 
     try:
-        async for redis in get_redis():
-            redis_key = f"{FAILED_DOMAINS_PREFIX}{domain}"
-            await redis.set(redis_key, domain, ex=expire_seconds)
-            logger.debug(f"已将失败域名 {domain} 存入Redis，过期时间：{expire_seconds}秒")
+        await set_cache(f"{domain}", domain, ttl=expire_seconds, prefix=FAILED_DOMAINS_PREFIX)
+
+        logger.debug(f"已将失败域名 {domain} 存入Redis，过期时间：{expire_seconds}秒")
     except Exception as e:
         logger.error(f"将失败域名存入Redis时出错：{e}")
 
 
 async def is_domain_failed(domain: str) -> bool:
-    """检查域名是否在Redis的失败列表中
-
-    Args:
-        domain: 要检查的域名
-
-    Returns:
-        True: 域名在失败列表中；False: 不在或Redis查询失败
-    """
     if not domain:
         return False
 
     try:
-        async for redis in get_redis():
-            redis_key = f"{FAILED_DOMAINS_PREFIX}{domain}"
-            result = await redis.exists(redis_key)
-            return result > 0
+        return await exist_cache(domain, prefix=FAILED_DOMAINS_PREFIX)
     except Exception as e:
         logger.error(f"检查域名是否失败时出错：{e}")
         return False
 
 
 async def delete_failed_domain(domain: str) -> None:
-    """从Redis中删除失败域名记录
-
-    Args:
-        domain: 要删除的域名
-    """
     if not domain:
         return
 
     try:
-        async for redis in get_redis():
-            redis_key = f"{FAILED_DOMAINS_PREFIX}{domain}"
-            await redis.delete(redis_key)
-            logger.debug(f"已从Redis删除失败域名 {domain}")
+        await remove_cache(domain, prefix=FAILED_DOMAINS_PREFIX)
+
+        logger.debug(f"已从Redis删除失败域名 {domain}")
     except Exception as e:
         logger.error(f"从Redis删除失败域名时出错：{e}")