diff --git a/aiosched/context.py b/aiosched/context.py index a13fd89..fd60263 100644 --- a/aiosched/context.py +++ b/aiosched/context.py @@ -58,6 +58,8 @@ class TaskScope: if not self.entry_point.done(): self.timed_out = True await throw(self.entry_point, TimeoutError("timed out")) + # Wait for the parent to exit or cancel us + await wait(self.entry_point) async def __aenter__(self): self.entry_point = await current_task() @@ -67,13 +69,13 @@ class TaskScope: return self async def __aexit__(self, exc_type: type, exception: Exception, tb): - await close_scope(self) if self.timeout and not self.waiter.done(): # Well, looks like we finished before our worker. # Thanks for your help! Now die. self.timeout_scope.cancellable = True await cancel(self.waiter, block=True) # Task scopes are sick: Nathaniel, you're an effing genius. + await close_scope(self) if isinstance(exception, TimeoutError) and self.timed_out: # This way we only silence our own timeouts and not # someone else's!