From 45bcc63c064c35efe0f28ccf8c84fc18a50e0281 Mon Sep 17 00:00:00 2001
From: jxxghp <jxxghp@gmail.com>
Date: Fri, 1 Aug 2025 11:48:37 +0800
Subject: [PATCH] fix rate_limit async

---
 app/utils/common.py | 43 +++++++++++++++++++++++++++++++++++++++++--
 app/utils/limit.py  | 37 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/app/utils/common.py b/app/utils/common.py
index 92d9cca6..44eea088 100644
--- a/app/utils/common.py
+++ b/app/utils/common.py
@@ -1,3 +1,4 @@
+import inspect
 import time
 from functools import wraps
 from typing import Any, Callable
@@ -34,7 +35,29 @@ def retry(ExceptionToCheck: Any,
                     mdelay *= backoff
             return f(*args, **kwargs)
 
-        return f_retry
+        async def async_f_retry(*args, **kwargs):
+            mtries, mdelay = tries, delay
+            while mtries > 1:
+                try:
+                    return await f(*args, **kwargs)
+                except ImmediateException:
+                    raise
+                except ExceptionToCheck as e:
+                    msg = f"{str(e)}, {mdelay} 秒后重试 ..."
+                    if logger:
+                        logger.warn(msg)
+                    else:
+                        print(msg)
+                    time.sleep(mdelay)
+                    mtries -= 1
+                    mdelay *= backoff
+            return await f(*args, **kwargs)
+
+        # 根据函数类型返回相应的包装器
+        if inspect.iscoroutinefunction(f):
+            return async_f_retry
+        else:
+            return f_retry
 
     return deco_retry
 
@@ -58,6 +81,22 @@ def log_execution_time(logger: Any = None):
                 print(msg)
             return result
 
-        return wrapper
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            start_time = time.time()
+            result = await func(*args, **kwargs)
+            end_time = time.time()
+            msg = f"{func.__name__} execution time: {end_time - start_time:.2f} seconds"
+            if logger:
+                logger.debug(msg)
+            else:
+                print(msg)
+            return result
+
+        # 根据函数类型返回相应的包装器
+        if inspect.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return wrapper
 
     return decorator
diff --git a/app/utils/limit.py b/app/utils/limit.py
index bec915da..6205d1fc 100644
--- a/app/utils/limit.py
+++ b/app/utils/limit.py
@@ -1,4 +1,5 @@
 import functools
+import inspect
 import threading
 import time
 from collections import deque
@@ -303,7 +304,41 @@ def rate_limit_handler(limiter: BaseRateLimiter, raise_on_limit: bool = False) -
                     raise e
                 return None
 
-        return wrapper
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs) -> Optional[Any]:
+            # 检查是否传入了 "raise_exception" 参数，优先使用该参数，否则使用默认的 raise_on_limit 值
+            raise_exception = kwargs.get("raise_exception", raise_on_limit)
+
+            # 检查是否可以进行调用，调用 limiter.can_call() 方法
+            can_call, message = limiter.can_call()
+            if not can_call:
+                # 如果调用受限，并且 raise_exception 为 True，则抛出限流异常
+                if raise_exception:
+                    raise RateLimitExceededException(message)
+                # 如果不抛出异常，则返回 None 表示跳过调用
+                return None
+
+            # 如果调用允许，执行目标函数，并记录一次调用
+            try:
+                result = await func(*args, **kwargs)
+                limiter.record_call()
+                if limiter.reset_on_success:
+                    limiter.reset()
+                return result
+            except LimitException as e:
+                # 如果目标函数触发了限流相关的异常，执行限流器的触发逻辑（如递增等待时间）
+                limiter.trigger_limit()
+                logger.error(limiter.format_log(f"触发限流：{str(e)}"))
+                # 如果 raise_exception 为 True，则抛出异常，否则返回 None
+                if raise_exception:
+                    raise e
+                return None
+
+        # 根据函数类型返回相应的包装器
+        if inspect.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return wrapper
 
     return decorator