32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
|
|
import asyncio
|
|
from functools import wraps
|
|
from pretor.utils.error import RetryableError
|
|
|
|
def retry_on_retryable_error(max_retries=3, base_delay=1):
|
|
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 as e:
|
|
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 as e:
|
|
if attempt == max_retries - 1:
|
|
raise
|
|
time.sleep(base_delay * (2 ** attempt))
|
|
return sync_wrapper
|
|
return decorator
|