diff options
author | Aldo Cortesi <aldo@corte.si> | 2018-04-17 10:19:55 +1200 |
---|---|---|
committer | Aldo Cortesi <aldo@corte.si> | 2018-04-17 10:42:40 +1200 |
commit | df83c0eef7befc9e9910806a12bbcd508023d5e1 (patch) | |
tree | 2779063f68318c176d4c4a00f85dc29087c2756a | |
parent | ef4db528871c96a0304a2057d5f1ee8330ca6122 (diff) | |
download | mitmproxy-df83c0eef7befc9e9910806a12bbcd508023d5e1.tar.gz mitmproxy-df83c0eef7befc9e9910806a12bbcd508023d5e1.tar.bz2 mitmproxy-df83c0eef7befc9e9910806a12bbcd508023d5e1.zip |
asyncio: fix exit behavior for console
- Add a master.run_loop function. This encapsulates our run behaviour so that
it can be used by implementations that need to manage their own run loop (like urwid).
- Shift crash exit message to the common core. I'm not convinced we really need
this, but if we want it it should be centralised.
- Clean up an extra exception that can be thrown by asyncio itself on "dirty"
termination after a mitmproxy crash.
-rw-r--r-- | mitmproxy/master.py | 32 | ||||
-rw-r--r-- | mitmproxy/tools/console/master.py | 21 |
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) |