asyncio: clarify shutdown semantics

This patch clarifies proxy shutdown, and specifies that the master.shutdown()
method is thread-save.
This commit is contained in:
Aldo Cortesi 2018-04-16 10:16:51 +12:00
parent 850c855495
commit 565146311a
6 changed files with 29 additions and 27 deletions

View File

@ -92,19 +92,27 @@ class Master:
try: try:
loop.run_forever() loop.run_forever()
finally: finally:
self.shutdown()
pending = asyncio.Task.all_tasks() pending = asyncio.Task.all_tasks()
loop.run_until_complete(asyncio.gather(*pending)) loop.run_until_complete(asyncio.gather(*pending))
loop.close() loop.close()
self.addons.trigger("done") self.addons.trigger("done")
async def _shutdown(self):
if self.server:
self.server.shutdown()
loop = asyncio.get_event_loop()
loop.stop()
def shutdown(self): def shutdown(self):
"""
Shut down the proxy. This method is thread-safe.
"""
if not self.should_exit.is_set(): if not self.should_exit.is_set():
if self.server:
self.server.shutdown()
self.should_exit.set() self.should_exit.set()
loop = asyncio.get_event_loop() asyncio.run_coroutine_threadsafe(
loop.stop() self._shutdown(),
loop = self.channel.loop,
)
def _change_reverse_host(self, f): def _change_reverse_host(self, f):
""" """

View File

@ -140,7 +140,7 @@ def mitmproxy(args=None): # pragma: no cover
assert_utf8_env() assert_utf8_env()
from mitmproxy.tools import console 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 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) m = run(dump.DumpMaster, cmdline.mitmdump, args, extra)
if m and m.errorcheck.has_errored: if m and m.errorcheck.has_errored:
sys.exit(1) sys.exit(1)
return m
def mitmweb(args=None): # pragma: no cover def mitmweb(args=None): # pragma: no cover
from mitmproxy.tools import web from mitmproxy.tools import web
run(web.master.WebMaster, cmdline.mitmweb, args) return run(web.master.WebMaster, cmdline.mitmweb, args)

View File

@ -1,9 +1,14 @@
import asyncio
import pytest
from mitmproxy.addons import keepserving from mitmproxy.addons import keepserving
from mitmproxy.test import taddons from mitmproxy.test import taddons
def test_keepserving(): @pytest.mark.asyncio
async def test_keepserving():
ks = keepserving.KeepServing() ks = keepserving.KeepServing()
with taddons.context(ks) as tctx: with taddons.context(ks) as tctx:
ks.event_processing_complete() ks.event_processing_complete()
asyncio.sleep(0.1)
assert tctx.master.should_exit.is_set() assert tctx.master.should_exit.is_set()

View File

@ -1,5 +0,0 @@
from mitmproxy import ctx
def tick():
ctx.master.shutdown()

View File

@ -1,25 +1,18 @@
import pytest import pytest
from mitmproxy.test import tutils
from mitmproxy.tools import main from mitmproxy.tools import main
shutdown_script = tutils.test_data.path("mitmproxy/data/addonscripts/shutdown.py")
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_mitmweb(): async def test_mitmweb(event_loop):
main.mitmweb([ m = main.mitmweb([
"--no-web-open-browser", "--no-web-open-browser",
"-q", "-q", "-p", "0",
"-p", "0",
"-s", shutdown_script
]) ])
await m._shutdown()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_mitmdump(): async def test_mitmdump():
main.mitmdump([ m = main.mitmdump(["-q", "-p", "0"])
"-q", await m._shutdown()
"-p", "0",
"-s", shutdown_script
])

View File

@ -40,7 +40,7 @@ class MasterTest:
async def dummy_cycle(self, master, n, content): async def dummy_cycle(self, master, n, content):
for i in range(n): for i in range(n):
await self.cycle(master, content) await self.cycle(master, content)
master.shutdown() await master._shutdown()
def flowfile(self, path): def flowfile(self, path):
with open(path, "wb") as f: with open(path, "wb") as f: