84 lines
2.5 KiB
Python
84 lines
2.5 KiB
Python
import asyncio
|
|
import concurrent.futures
|
|
import functools
|
|
import inspect
|
|
import sys
|
|
import time
|
|
import typing
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any, AsyncGenerator, Callable, Iterator, List, Optional, Set, TypeVar, cast
|
|
|
|
import synchronicity
|
|
from typing_extensions import ParamSpec
|
|
|
|
|
|
synchronizer = synchronicity.Synchronizer()
|
|
|
|
def synchronize_api(obj, target_module=None):
|
|
if inspect.isclass(obj):
|
|
blocking_name = obj.__name__.lstrip("_")
|
|
elif inspect.isfunction(object):
|
|
blocking_name = obj.__name__.lstrip("_")
|
|
elif isinstance(obj, TypeVar):
|
|
blocking_name = "_BLOCKING_" + obj.__name__
|
|
else:
|
|
blocking_name = None
|
|
if target_module is None:
|
|
target_module = obj.__module__
|
|
return synchronizer.create_blocking(obj, blocking_name, target_module=target_module)
|
|
|
|
def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout=90):
|
|
"""Decorator that calls an async function multiple times, with a given timeout.
|
|
|
|
If a `base_delay` is provided, the function is given an exponentially
|
|
increasing delay on each run, up until the maximum number of attempts.
|
|
|
|
Usage:
|
|
|
|
```
|
|
@retry
|
|
async def may_fail_default():
|
|
# ...
|
|
pass
|
|
|
|
@retry(n_attempts=5, base_delay=1)
|
|
async def may_fail_delay():
|
|
# ...
|
|
pass
|
|
```
|
|
"""
|
|
|
|
def decorator(fn):
|
|
@functools.wraps(fn)
|
|
async def f_wrapped(*args, **kwargs):
|
|
delay = base_delay
|
|
for i in range(n_attempts):
|
|
t0 = time.time()
|
|
try:
|
|
return await asyncio.wait_for(fn(*args, **kwargs), timeout=timeout)
|
|
except asyncio.CancelledError:
|
|
#logger.debug(f"Function {fn} was cancelled")
|
|
raise
|
|
except Exception as e:
|
|
if i >= n_attempts - 1:
|
|
raise
|
|
"""
|
|
logger.debug(
|
|
f"Failed invoking function {fn}: {e}"
|
|
f" (took {time.time() - t0}s, sleeping {delay}s"
|
|
f" and trying {n_attempts - i - 1} more times)"
|
|
)
|
|
"""
|
|
await asyncio.sleep(delay)
|
|
delay *= delay_factor
|
|
|
|
return f_wrapped
|
|
|
|
if direct_fn is not None:
|
|
# It's invoked like @retry
|
|
return decorator(direct_fn)
|
|
else:
|
|
# It's invoked like @retry(n_attempts=...)
|
|
return decorator
|
|
|