From 23728c27133b7a4cb64f135e9961b0985d8a4134 Mon Sep 17 00:00:00 2001
From: Jan Pochyla <jpochyla@gmail.com>
Date: Thu, 18 Jan 2018 14:38:56 +0100
Subject: [PATCH] loop: make wait properly dispose of child tasks

---
 src/trezor/loop.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/trezor/loop.py b/src/trezor/loop.py
index bfaffa0fd..91abe96c9 100644
--- a/src/trezor/loop.py
+++ b/src/trezor/loop.py
@@ -251,8 +251,8 @@ class wait(Syscall):
         self.children = children
         self.wait_for = wait_for
         self.exit_others = exit_others
-        self.scheduled = None
-        self.finished = None
+        self.scheduled = None  # list of scheduled wrapper tasks
+        self.finished = None  # list of children that finished
         self.callback = None
 
     def handle(self, task):
@@ -263,9 +263,8 @@ class wait(Syscall):
             schedule(ct)
 
     def exit(self):
-        for task in self.scheduled:
-            if task not in self.finished:
-                close(task)
+        for ct in self.scheduled:
+            close(ct)
 
     async def _wait(self, child):
         try:
@@ -278,14 +277,14 @@ class wait(Syscall):
     def _finish(self, child, result):
         self.finished.append(child)
         if self.wait_for == len(self.finished) or isinstance(result, Exception):
+            schedule(self.callback, result)
             if self.exit_others:
                 self.exit()
-            schedule(self.callback, result)
 
     def __iter__(self):
         try:
             return (yield self)
-        except Exception:
+        except:
             # exception was raised on the waiting task externally with
             # close() or throw(), kill the children tasks and re-raise
             self.exit()