From 565146311a5f7939f107af12a91d94f5f96d56fc Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@corte.si>
Date: Mon, 16 Apr 2018 10:16:51 +1200
Subject: [PATCH] asyncio: clarify shutdown semantics

This patch clarifies proxy shutdown, and specifies that the master.shutdown()
method is thread-save.
---
 mitmproxy/master.py                          | 18 +++++++++++++-----
 mitmproxy/tools/main.py                      |  5 +++--
 test/mitmproxy/addons/test_keepserving.py    |  7 ++++++-
 test/mitmproxy/data/addonscripts/shutdown.py |  5 -----
 test/mitmproxy/tools/test_main.py            | 19 ++++++-------------
 test/mitmproxy/tservers.py                   |  2 +-
 6 files changed, 29 insertions(+), 27 deletions(-)
 delete mode 100644 test/mitmproxy/data/addonscripts/shutdown.py

diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index dff8ca5c9..19a2ac672 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -92,19 +92,27 @@ class Master:
         try:
             loop.run_forever()
         finally:
-            self.shutdown()
             pending = asyncio.Task.all_tasks()
             loop.run_until_complete(asyncio.gather(*pending))
             loop.close()
         self.addons.trigger("done")
 
+    async def _shutdown(self):
+        if self.server:
+            self.server.shutdown()
+        loop = asyncio.get_event_loop()
+        loop.stop()
+
     def shutdown(self):
+        """
+            Shut down the proxy. This method is thread-safe.
+        """
         if not self.should_exit.is_set():
-            if self.server:
-                self.server.shutdown()
             self.should_exit.set()
-            loop = asyncio.get_event_loop()
-            loop.stop()
+            asyncio.run_coroutine_threadsafe(
+                self._shutdown(),
+                loop = self.channel.loop,
+            )
 
     def _change_reverse_host(self, f):
         """
diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py
index e8ec993dd..8d63aa323 100644
--- a/mitmproxy/tools/main.py
+++ b/mitmproxy/tools/main.py
@@ -140,7 +140,7 @@ def mitmproxy(args=None):  # pragma: no cover
     assert_utf8_env()
 
     from mitmproxy.tools import console
-    run(console.master.ConsoleMaster, cmdline.mitmproxy, args)
+    return run(console.master.ConsoleMaster, cmdline.mitmproxy, args)
 
 
 def mitmdump(args=None):  # pragma: no cover
@@ -159,8 +159,9 @@ def mitmdump(args=None):  # pragma: no cover
     m = run(dump.DumpMaster, cmdline.mitmdump, args, extra)
     if m and m.errorcheck.has_errored:
         sys.exit(1)
+    return m
 
 
 def mitmweb(args=None):  # pragma: no cover
     from mitmproxy.tools import web
-    run(web.master.WebMaster, cmdline.mitmweb, args)
+    return run(web.master.WebMaster, cmdline.mitmweb, args)
diff --git a/test/mitmproxy/addons/test_keepserving.py b/test/mitmproxy/addons/test_keepserving.py
index 2869d0976..5eafa7929 100644
--- a/test/mitmproxy/addons/test_keepserving.py
+++ b/test/mitmproxy/addons/test_keepserving.py
@@ -1,9 +1,14 @@
+import asyncio
+import pytest
+
 from mitmproxy.addons import keepserving
 from mitmproxy.test import taddons
 
 
-def test_keepserving():
+@pytest.mark.asyncio
+async def test_keepserving():
     ks = keepserving.KeepServing()
     with taddons.context(ks) as tctx:
         ks.event_processing_complete()
+        asyncio.sleep(0.1)
         assert tctx.master.should_exit.is_set()
diff --git a/test/mitmproxy/data/addonscripts/shutdown.py b/test/mitmproxy/data/addonscripts/shutdown.py
deleted file mode 100644
index 3da4d03eb..000000000
--- a/test/mitmproxy/data/addonscripts/shutdown.py
+++ /dev/null
@@ -1,5 +0,0 @@
-from mitmproxy import ctx
-
-
-def tick():
-    ctx.master.shutdown()
diff --git a/test/mitmproxy/tools/test_main.py b/test/mitmproxy/tools/test_main.py
index a514df740..ba2137331 100644
--- a/test/mitmproxy/tools/test_main.py
+++ b/test/mitmproxy/tools/test_main.py
@@ -1,25 +1,18 @@
 import pytest
 
-from mitmproxy.test import tutils
 from mitmproxy.tools import main
 
-shutdown_script = tutils.test_data.path("mitmproxy/data/addonscripts/shutdown.py")
-
 
 @pytest.mark.asyncio
-async def test_mitmweb():
-    main.mitmweb([
+async def test_mitmweb(event_loop):
+    m = main.mitmweb([
         "--no-web-open-browser",
-        "-q",
-        "-p", "0",
-        "-s", shutdown_script
+        "-q", "-p", "0",
     ])
+    await m._shutdown()
 
 
 @pytest.mark.asyncio
 async def test_mitmdump():
-    main.mitmdump([
-        "-q",
-        "-p", "0",
-        "-s", shutdown_script
-    ])
+    m = main.mitmdump(["-q", "-p", "0"])
+    await m._shutdown()
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index 13ee4a430..857ca45d0 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -40,7 +40,7 @@ class MasterTest:
     async def dummy_cycle(self, master, n, content):
         for i in range(n):
             await self.cycle(master, content)
-        master.shutdown()
+        await master._shutdown()
 
     def flowfile(self, path):
         with open(path, "wb") as f: