Closing all external coroutines when the scheduler shuts down or the …

…individual task finishes.

mango/util/scheduling.py

@@ -1,6 +1,7 @@
 """
 Module for commonly used time based scheduled task executed inside one agent.
 """
+
 import asyncio
 import concurrent.futures
 import datetime
@@ -91,6 +92,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
@@ -131,6 +139,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):
@@ -155,6 +168,9 @@ async def run(self):
         await self._wait(self._timestamp)
         return await self._coro
 
+    def close(self):
+        _close_coro(self._coro)
+
 
 class AwaitingTask(ScheduledTask):
     """
@@ -175,6 +191,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):
     """
@@ -188,6 +208,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):
     """
@@ -274,6 +297,9 @@ async def run(self):
             self.notify_running()
         return await self._coro
 
+    def close(self):
+        _close_coro(self._coro)
+
 
 # process tasks
 
@@ -838,4 +864,6 @@ def shutdown(self):
         for _, _, event, _ in self._scheduled_process_tasks:
             if event is not None:
                 event.set()
+        for task, _, _, _ in self._scheduled_tasks:
+            task.close()
         self._process_pool_exec.shutdown()

tests/unit_tests/util/scheduling_test.py

@@ -34,7 +34,9 @@ async def increase_counter():
     )
 
     # WHEN
-    t = scheduler.schedule_task(RecurrentScheduledTask(increase_counter, recurrency, clock))
+    t = scheduler.schedule_task(
+        RecurrentScheduledTask(increase_counter, recurrency, clock)
+    )
     try:
         new_time = start + datetime.timedelta(days=1)
         clock.set_time(new_time.timestamp())
@@ -73,6 +75,7 @@ async def increase_counter():
     assert task._is_done.done()
     assert len(l) == 2
 
+
 @pytest.mark.asyncio
 async def test_recurrent_wait():
     # GIVEN
@@ -84,14 +87,15 @@ async def test_recurrent_wait():
 
     async def increase_counter():
         l.append(clock._time)
+
     tomorrow = start + datetime.timedelta(days=1)
     aftertomorrow = start + datetime.timedelta(days=2)
-    recurrency = rrule.rrule(
-        rrule.DAILY, interval=1, dtstart=tomorrow, until=end
-    )
+    recurrency = rrule.rrule(rrule.DAILY, interval=1, dtstart=tomorrow, until=end)
 
     # WHEN
-    t = scheduler.schedule_task(RecurrentScheduledTask(increase_counter, recurrency, clock))
+    t = scheduler.schedule_task(
+        RecurrentScheduledTask(increase_counter, recurrency, clock)
+    )
     task = scheduler._scheduled_tasks[0][0]
     try:
         clock.set_time(start.timestamp())
@@ -175,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