Smart cancel of failed process tasks.

mango/agent/core.py

@@ -472,7 +472,7 @@ async def shutdown(self):
         except asyncio.CancelledError:
             pass
         try:
-            self._scheduler.shutdown()
+            await self._scheduler.shutdown()
         except asyncio.CancelledError:
             pass
         finally:

mango/util/scheduling.py

@@ -810,6 +810,8 @@ def _remove_process_task(self, fut=asyncio.Future):
         for i in range(len(self._scheduled_process_tasks)):
             _, task, event, _ = self._scheduled_process_tasks[i]
             if task == fut:
+                event[0].set()
+                event[1].set()
                 del self._scheduled_process_tasks[i]
                 break
 
@@ -869,7 +871,7 @@ async def tasks_complete_or_sleeping(self):
                     # we need to recognize how many sleeping tasks we have in order to find out if all tasks are done
                     sleeping_tasks.append(scheduled_task)
 
-    def shutdown(self):
+    async def shutdown(self):
         """
         Shutdown internal process executor pool.
         """
@@ -879,5 +881,6 @@ def shutdown(self):
                 event[1].set()
         for task, _, _, _ in self._scheduled_tasks:
             task.close()
+        await self.stop()
         if self._process_pool_exec is not None:
             self._process_pool_exec.shutdown()

tests/unit_tests/util/scheduling_test.py

@@ -179,7 +179,7 @@ async def increase_counter():
     with pytest.raises(asyncio.exceptions.TimeoutError):
         await asyncio.wait_for(t, timeout=0.2)
 
-    scheduler.shutdown()
+    await scheduler.shutdown()
 
     # THEN
     assert len(l) == 0
@@ -368,7 +368,9 @@ async def test_task_as_process_suspend():
     except asyncio.exceptions.TimeoutError as err:
         pass
 
-    scheduler.shutdown()
+    scheduler.resume(marker)
+
+    await scheduler.shutdown()
 
 
 @pytest.mark.asyncio