Fix resume/pause

asynq/contexts.py

@@ -20,7 +20,8 @@
 from ._debug import options as _debug_options
 
 
-ASYNCIO_CONTEXT_FIELD = '_asynq_contexts'
+ASYNCIO_CONTEXT_FIELD = "_asynq_contexts"
+ASYNCIO_CONTEXT_ACTIVE_FIELD = "_asynq_contexts_active"
 
 
 class NonAsyncContext(object):
@@ -73,6 +74,7 @@ def leave_context(context, active_task):
     if active_task is not None:
         active_task._leave_context(context)
 
+
 def enter_context_asyncio(context):
     if _debug_options.DUMP_CONTEXTS:
         debug.write("@async: +context: %s" % debug.str(context))
@@ -85,21 +87,27 @@ def enter_context_asyncio(context):
     else:
         setattr(task, ASYNCIO_CONTEXT_FIELD, {id(context): context})
 
+
 def leave_context_asyncio(context):
     if _debug_options.DUMP_CONTEXTS:
         debug.write("@async: -context: %s" % debug.str(context))
 
     task = asyncio.current_task()
     del getattr(task, ASYNCIO_CONTEXT_FIELD)[id(context)]  # type: ignore
-    context.pause()
+
 
 def pause_contexts_asyncio(task):
-    for ctx in reversed(list(getattr(task, ASYNCIO_CONTEXT_FIELD, {}).values())):
-        ctx.pause()
+    if getattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, False):
+        setattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, False)
+        for ctx in reversed(list(getattr(task, ASYNCIO_CONTEXT_FIELD, {}).values())):
+            ctx.pause()
+
 
 def resume_contexts_asyncio(task):
-    for ctx in getattr(task, ASYNCIO_CONTEXT_FIELD, {}).values():
-        ctx.resume()
+    if not getattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, False):
+        setattr(task, ASYNCIO_CONTEXT_ACTIVE_FIELD, True)
+        for ctx in getattr(task, ASYNCIO_CONTEXT_FIELD, {}).values():
+            ctx.resume()
 
 
 class AsyncContext(object):

asynq/decorators.py

@@ -21,13 +21,15 @@
 import qcore.helpers as core_helpers
 import qcore.inspection as core_inspection
 
-from . import _debug, async_task, asynq_to_async, futures
-from .contexts import ASYNCIO_CONTEXT_FIELD, pause_contexts_asyncio, resume_contexts_asyncio
+from . import async_task, asynq_to_async, futures
+from .contexts import (
+    ASYNCIO_CONTEXT_FIELD,
+    pause_contexts_asyncio,
+    resume_contexts_asyncio,
+)
 
 __traceback_hide__ = True
 
-_debug_options = _debug.options
-
 
 def lazy(fn):
     """Converts a function into a lazy one - i.e. its call
@@ -175,15 +177,14 @@ async def wrapped(*_args, **_kwargs):
                         send = None
                         generator = self.fn(*_args, **_kwargs)
                         while True:
+                            resume_contexts_asyncio(task)
                             try:
                                 result = generator.send(send)
                             except StopIteration as exc:
                                 return exc.value
 
-                            # pause the current task's contexts
                             pause_contexts_asyncio(task)
                             send = await asynq_to_async.resolve_awaitables(result)
-                            resume_contexts_asyncio(task)
 
                 self.asyncio_fn = wrapped
             else:

asynq/tests/test_asynq_to_async.py

@@ -19,58 +19,39 @@
 
 import asynq
 from asynq.tools import AsyncTimer
-from asynq.batching import BatchItemBase, BatchBase
 
 
-async def f3():
-    return 200
-
-
-@asynq.asynq(asyncio_fn=f3)
-def f2():
-    return 100
-
-
-@asynq.asynq()
-def f(x):
-    a = yield asynq.ConstFuture(3)
-    b = yield asynq.ConstFuture(2)
-    assert (yield None) == None
-    return a - b + x
+def test_asyncio():
+    async def f3():
+        return 200
 
+    @asynq.asynq(asyncio_fn=f3)
+    def f2():
+        return 100
 
-@asynq.asynq()
-def g(x):
-    obj = yield {
-        "a": [f.asynq(0), f.asynq(1)],
-        "b": (f.asynq(2), f.asynq(3)),
-        "c": f.asynq(4),
-        "d": f2.asynq(),
-    }
-    return obj
+    @asynq.asynq()
+    def f(x):
+        a = yield asynq.ConstFuture(3)
+        b = yield asynq.ConstFuture(2)
+        assert (yield None) == None
+        return a - b + x
 
+    @asynq.asynq()
+    def g(x):
+        obj = yield {
+            "a": [f.asynq(0), f.asynq(1)],
+            "b": (f.asynq(2), f.asynq(3)),
+            "c": f.asynq(4),
+            "d": f2.asynq(),
+        }
+        return obj
 
-def test_asyncio():
     assert asyncio.run(g.asyncio(5)) == {"a": [1, 2], "b": (3, 4), "c": 5, "d": 200}
 
 
-
-
 def test_context():
-    class TestBatchItem(BatchItemBase):
-        pass
-
-    class TestBatch(BatchBase):
-        def _try_switch_active_batch(self):
-            pass
-        def _flush(self):
-            for item in self.items:
-                item.set_value(None)
-            time.sleep(0.1 * len(self.items))
-        def _cancel(self):
-            pass
-
-    batch = TestBatch()
+    async def blocking_op():
+        await asyncio.sleep(0.1)
 
     @asynq.asynq()
     def f1():
@@ -87,16 +68,16 @@ def f2():
             t = yield f3.asynq()
             time.sleep(0.1)
         return timer.total_time, t
-    
+
     @asynq.asynq()
     def f3():
         with AsyncTimer() as timer:
             # since AsyncTimer is paused on blocking operations,
             # the time for TestBatch is not measured
-            yield [TestBatchItem(batch), TestBatchItem(batch)]
+            yield [blocking_op(), blocking_op()]
         return timer.total_time
 
-    t1, t2, t3 = f1()
+    t1, t2, t3 = asyncio.run(f1.asyncio())
     assert_eq(400000, t1, tolerance=10000)  # 400ms, 10us tolerance
     assert_eq(200000, t2, tolerance=10000)  # 200ms, 10us tolerance
     assert_eq(000000, t3, tolerance=10000)  #   0ms, 10us tolerance

asynq/tools.py

@@ -453,9 +453,13 @@ def __init__(self):
 
     def resume(self):
         self._last_start_time = utime()
+        if hasattr(self, "debug"):
+            print("resume")
 
     def pause(self):
         self.total_time += utime() - self._last_start_time
+        if hasattr(self, "debug"):
+            print("pause - ", self.total_time)
 
 
 class AsyncEventHook(EventHook):