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:
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):
"""

View File

@ -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)

View File

@ -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()

View File

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

View File

@ -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()

View File

@ -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: