Merge dev

mango/util/scheduling.py

@@ -96,6 +96,13 @@ def coro(self):
 # asyncio tasks
 
 
+def _close_coro(coro):
+    try:
+        coro.close()
+    except:
+        pass
+
+
 class ScheduledTask:
     """
     Base class for scheduled tasks in mango. Within this class it is possible to
@@ -136,6 +143,11 @@ def on_stop(self, fut: asyncio.Future = None):
             self._on_stop_hook_in(fut)
         if self._is_observable:
             self._is_done.set_result(True)
+        self.close()
+
+    def close(self):
+        """Perform closing actions"""
+        pass
 
 
 class TimestampScheduledTask(ScheduledTask):
@@ -160,6 +172,9 @@ async def run(self):
         await self._wait(self._timestamp)
         return await self._coro
 
+    def close(self):
+        _close_coro(self._coro)
+
 
 class AwaitingTask(ScheduledTask):
     """
@@ -180,6 +195,10 @@ async def run(self):
         self.notify_running()
         return await self._coroutine
 
+    def close(self):
+        _close_coro(self._awaited_coroutine)
+        _close_coro(self._coroutine)
+
 
 class InstantScheduledTask(TimestampScheduledTask):
     """
@@ -193,6 +212,9 @@ def __init__(self, coroutine, clock: Clock = None, on_stop=None, observable=True
             coroutine, clock.time, clock=clock, on_stop=on_stop, observable=observable
         )
 
+    def close(self):
+        _close_coro(self._coro)
+
 
 class PeriodicScheduledTask(ScheduledTask):
     """
@@ -279,6 +301,9 @@ async def run(self):
             self.notify_running()
         return await self._coro
 
+    def close(self):
+        _close_coro(self._coro)
+
 
 # process tasks
 
@@ -852,6 +877,8 @@ def shutdown(self):
         for _, _, event, _ in self._scheduled_process_tasks:
             if event[1] is not None:
                 event[1].set()
+        for task, _, _, _ in self._scheduled_tasks:
+            task.close()
         self._process_pool_exec.shutdown(wait=True, cancel_futures=True)
         if self._manager is not None:
             self._manager.shutdown()

tests/unit_tests/util/scheduling_test.py

@@ -179,6 +179,8 @@ async def increase_counter():
     with pytest.raises(asyncio.exceptions.TimeoutError):
         await asyncio.wait_for(t, timeout=0.2)
 
+    scheduler.shutdown()
+
     # THEN
     assert len(l) == 0