retry.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. from asyncio import sleep
  2. from typing import (
  3. TYPE_CHECKING,
  4. Any,
  5. Awaitable,
  6. Callable,
  7. Optional,
  8. Tuple,
  9. Type,
  10. TypeVar,
  11. Union,
  12. )
  13. from redis.exceptions import ConnectionError, RedisError, TimeoutError
  14. from redis.retry import AbstractRetry
  15. T = TypeVar("T")
  16. if TYPE_CHECKING:
  17. from redis.backoff import AbstractBackoff
  18. class Retry(AbstractRetry[RedisError]):
  19. __hash__ = AbstractRetry.__hash__
  20. def __init__(
  21. self,
  22. backoff: "AbstractBackoff",
  23. retries: int,
  24. supported_errors: Tuple[Type[RedisError], ...] = (
  25. ConnectionError,
  26. TimeoutError,
  27. ),
  28. ):
  29. super().__init__(backoff, retries, supported_errors)
  30. def __eq__(self, other: Any) -> bool:
  31. if not isinstance(other, Retry):
  32. return NotImplemented
  33. return (
  34. self._backoff == other._backoff
  35. and self._retries == other._retries
  36. and set(self._supported_errors) == set(other._supported_errors)
  37. )
  38. async def call_with_retry(
  39. self,
  40. do: Callable[[], Awaitable[T]],
  41. fail: Union[
  42. Callable[[Exception], Any],
  43. Callable[[Exception, int], Any],
  44. ],
  45. is_retryable: Optional[Callable[[Exception], bool]] = None,
  46. with_failure_count: bool = False,
  47. ) -> T:
  48. """
  49. Execute an operation that might fail and returns its result, or
  50. raise the exception that was thrown depending on the `Backoff` object.
  51. `do`: the operation to call. Expects no argument.
  52. `fail`: the failure handler, expects the last error that was thrown
  53. ``is_retryable``: optional function to determine if an error is retryable
  54. ``with_failure_count``: if True, the failure count is passed to the failure handler
  55. """
  56. self._backoff.reset()
  57. failures = 0
  58. while True:
  59. try:
  60. return await do()
  61. except self._supported_errors as error:
  62. if is_retryable and not is_retryable(error):
  63. raise
  64. failures += 1
  65. if with_failure_count:
  66. await fail(error, failures)
  67. else:
  68. await fail(error)
  69. if self._retries >= 0 and failures > self._retries:
  70. raise error
  71. backoff = self._backoff.compute(failures)
  72. if backoff > 0:
  73. await sleep(backoff)