62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
# Copyright 2026 zhaoxi826
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
|
|
import asyncio
|
|
from functools import wraps
|
|
from kilostar.utils.error import RetryableError
|
|
|
|
|
|
def retry_on_retryable_error(max_retries=3, base_delay=1):
|
|
"""指数退避重试装饰器：仅在抛出 ``RetryableError`` 时重试。
|
|
|
|
同步/异步函数都支持。第 n 次重试前会 ``sleep(base_delay * 2**n)``。
|
|
|
|
Args:
|
|
max_retries: 最多尝试次数（含首次），超过后会把最后一次异常重新抛出。
|
|
base_delay: 退避基准秒数。
|
|
"""
|
|
|
|
def decorator(func):
|
|
if asyncio.iscoroutinefunction(func):
|
|
|
|
@wraps(func)
|
|
async def async_wrapper(*args, **kwargs):
|
|
for attempt in range(max_retries):
|
|
try:
|
|
return await func(*args, **kwargs)
|
|
except RetryableError:
|
|
if attempt == max_retries - 1:
|
|
raise
|
|
await asyncio.sleep(base_delay * (2**attempt))
|
|
|
|
return async_wrapper
|
|
else:
|
|
|
|
@wraps(func)
|
|
def sync_wrapper(*args, **kwargs):
|
|
import time
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
return func(*args, **kwargs)
|
|
except RetryableError:
|
|
if attempt == max_retries - 1:
|
|
raise
|
|
time.sleep(base_delay * (2**attempt))
|
|
|
|
return sync_wrapper
|
|
|
|
return decorator
|