Merge pull request #3064 from cortesi/mexit

asyncio: fix exit behaviour for console

2 files changed, 29 insertions, 24 deletions

diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 19a2ac67..bbbd07d0 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -1,3 +1,5 @@
+import sys
+import traceback
 import threading
 import asyncio
 import logging
@@ -85,18 +87,40 @@ class Master:
             self.addons.trigger("tick")
             await asyncio.sleep(0.1)
 
-    def run(self):
+    def run_loop(self, loop):
         self.start()
         asyncio.ensure_future(self.tick())
-        loop = asyncio.get_event_loop()
+
+        exc = None
         try:
-            loop.run_forever()
+            loop()
+        except Exception as e:
+            exc = traceback.format_exc()
         finally:
+            if not self.should_exit.is_set():
+                self.shutdown()
             pending = asyncio.Task.all_tasks()
-            loop.run_until_complete(asyncio.gather(*pending))
+            loop = asyncio.get_event_loop()
+            try:
+                loop.run_until_complete(asyncio.gather(*pending))
+            except Exception as e:
+                # When we exit with an error, shutdown might not happen cleanly,
+                # and we can get exceptions here caused by pending Futures.
+                pass
             loop.close()
+
+        if exc:  # pragma: no cover
+            print(exc, file=sys.stderr)
+            print("mitmproxy has crashed!", file=sys.stderr)
+            print("Please lodge a bug report at:", file=sys.stderr)
+            print("\thttps://github.com/mitmproxy/mitmproxy", file=sys.stderr)
+
         self.addons.trigger("done")
 
+    def run(self, func=None):
+        loop = asyncio.get_event_loop()
+        self.run_loop(loop.run_forever)
+
     async def _shutdown(self):
         if self.server:
             self.server.shutdown()
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index c66129b2..9ed73ce9 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -9,7 +9,6 @@ import stat
 import subprocess
 import sys
 import tempfile
-import traceback
 import typing  # noqa
 import contextlib
 
@@ -205,7 +204,6 @@ class ConsoleMaster(master.Master):
             screen = self.ui,
             handle_mouse = self.options.console_mouse,
         )
-
         self.window = window.Window(self)
         self.loop.widget = self.window
         self.window.refresh()
@@ -216,24 +214,7 @@ class ConsoleMaster(master.Master):
                 self.start_err = None
             self.loop.set_alarm_in(0.01, display_err)
 
-        self.start()
-        try:
-            self.loop.run()
-        except Exception:
-            self.loop.stop()
-            sys.stdout.flush()
-            print(traceback.format_exc(), file=sys.stderr)
-            print("mitmproxy has crashed!", file=sys.stderr)
-            print("Please lodge a bug report at:", file=sys.stderr)
-            print("\thttps://github.com/mitmproxy/mitmproxy", file=sys.stderr)
-            print("Shutting down...", file=sys.stderr)
-        finally:
-            sys.stderr.flush()
-            super().shutdown()
-        self.addons.trigger("done")
-
-    def shutdown(self):
-        raise urwid.ExitMainLoop
+        super().run_loop(self.loop.run)
 
     def overlay(self, widget, **kwargs):
         self.window.set_overlay(widget, **kwargs)