asynq/tools: update aretry to compatible with asynq_to_asyncio

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