From 9d92ef4b0c5cd89a7bad4a66bbfa2d53d098d3fc Mon Sep 17 00:00:00 2001 From: Joungjin Lee Date: Thu, 12 Dec 2024 01:00:43 -0800 Subject: [PATCH] asynq/tools: update aretry to compatible with asynq_to_asyncio --- asynq/tools.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/asynq/tools.py b/asynq/tools.py index 10216ac..c81dde9 100644 --- a/asynq/tools.py +++ b/asynq/tools.py @@ -21,6 +21,7 @@ import functools import inspect import itertools +import asyncio import threading import time import weakref @@ -282,6 +283,12 @@ def dirty(): return decorator +async def _asyncio_sleep(secs: float) -> None: + await asyncio.sleep(secs) + +@asynq(asyncio_fn=_asyncio_sleep) +def _asynq_sleep(secs: float) -> None: + time.sleep(secs) def aretry(exception_cls, max_tries=10, sleep=0.05): """Decorator for retrying an async function if it throws an exception. @@ -290,6 +297,8 @@ def aretry(exception_cls, max_tries=10, sleep=0.05): max_tries - maximum number of times this function can be executed sleep - number of seconds to sleep between function retries + Compatible with asyncio mode. + """ assert max_tries > 0, "max_tries (%d) should be a positive integer" % max_tries @@ -304,7 +313,7 @@ def wrapper(*args, **kwargs): except exception_cls: if i + 1 == max_tries: raise - time.sleep(sleep) + yield _asynq_sleep.asynq(sleep) # so that qcore.inspection.get_original_fn can retrieve the original function wrapper.original_fn = fn