From 499ea4da85c83a416f1f7461e5c522df8ce9a460 Mon Sep 17 00:00:00 2001
From: nocturn9x <hackhab@gmail.com>
Date: Sat, 4 Apr 2020 17:14:04 +0000
Subject: [PATCH] Fixed the TaskManager exception handling

---
 experiment.py   |  2 +-
 giambio/core.py | 12 ++++++++----
 giambio/util.py |  2 ++
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/experiment.py b/experiment.py
index b918202..0d6fb99 100644
--- a/experiment.py
+++ b/experiment.py
@@ -30,7 +30,7 @@ async def countdown(n):
         return "Count DOWN over"
     except giambio.CancelledError:
         print("countdown cancelled!")
-        raise Exception("Oh no!")   #TODO Propagate this
+        raise Exception("Oh no!")   # TODO Propagate this
 
 async def count(stop, step=1):
     try:
diff --git a/giambio/core.py b/giambio/core.py
index 187b79d..bce54f8 100644
--- a/giambio/core.py
+++ b/giambio/core.py
@@ -4,7 +4,7 @@ from collections import deque, defaultdict
 from selectors import DefaultSelector, EVENT_READ, EVENT_WRITE
 from heapq import heappush, heappop
 import socket
-from .exceptions import AlreadyJoinedError, CancelledError
+from .exceptions import AlreadyJoinedError, CancelledError, GiambioError
 from timeit import default_timer
 from time import sleep as wait
 from .socket import AsyncSocket, WantRead, WantWrite
@@ -29,6 +29,7 @@ class EventLoop:
         self.joined = defaultdict(list)  # Tasks that want to join
         self.clock = default_timer  # Monotonic clock to keep track of elapsed time
         self.sequence = 0  # To avoid TypeError in the (unlikely) event of two task with the same deadline we use a unique and incremental integer pushed to the queue together with the deadline and the function itself
+        self._exiting = False
 
     def loop(self):
         """Main event loop for giambio"""
@@ -63,9 +64,12 @@ class EventLoop:
                     self.running.execution = "CANCELLED"
                     self.to_run.extend(self.joined.pop(self.running, ()))
                 except Exception as err:
-                    self.running.execution = "ERRORED"
-                    self.running.result = Result(None, err)
-                    self.to_run.extend(self.joined.pop(self.running, ()))   # Reschedules the parent task
+                    if not self._exiting:
+                        self.running.execution = "ERRORED"
+                        self.running.result = Result(None, err)
+                        self.to_run.extend(self.joined.pop(self.running, ()))   # Reschedules the parent task
+                    else:
+                        raise
                 except KeyboardInterrupt:
                     self.running.throw(KeyboardInterrupt)
 
diff --git a/giambio/util.py b/giambio/util.py
index 4c121f1..33d96ec 100644
--- a/giambio/util.py
+++ b/giambio/util.py
@@ -17,8 +17,10 @@ class TaskManager:
 
 
     async def _cancel_and_raise(self, exc):
+        self.loop._exiting = True   # Tells the loop not to catch all exceptions
         for task in self.tasks:
             await task.cancel()
+        raise exc
 
     async def __aenter__(self):
         return self