asyncio: shift script reloading out of the tick event

The tick event is a nasty compromise, left over from when we didn't have an
event loop. This is the first patch in a series that explores moving our
built-in addons to managing coroutines on the eventloop directly for periodic
tasks.

13 files changed, 133 insertions, 105 deletions

diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py
index 9d0e7069..ac561073 100644
--- a/mitmproxy/addonmanager.py
+++ b/mitmproxy/addonmanager.py
@@ -140,6 +140,7 @@ class AddonManager:
         """
         for i in self.chain:
             self.remove(i)
+        self.lookup = {}
 
     def get(self, name):
         """
diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py
index 51341f2f..decd0759 100644
--- a/mitmproxy/addons/script.py
+++ b/mitmproxy/addons/script.py
@@ -1,7 +1,7 @@
+import asyncio
 import os
 import importlib.util
 import importlib.machinery
-import time
 import sys
 import types
 import typing
@@ -59,13 +59,15 @@ def script_error_handler(path, exc, msg="", tb=False):
     ctx.log.error(log_msg)
 
 
+ReloadInterval = 1
+
+
 class Script:
     """
         An addon that manages a single script.
     """
-    ReloadInterval = 2
 
-    def __init__(self, path):
+    def __init__(self, path: str, reload: bool) -> None:
         self.name = "scriptmanager:" + path
         self.path = path
         self.fullpath = os.path.expanduser(
@@ -73,45 +75,57 @@ class Script:
         )
         self.ns = None
 
-        self.last_load = 0
-        self.last_mtime = 0
         if not os.path.isfile(self.fullpath):
             raise exceptions.OptionsError('No such script')
 
+        self.reloadtask = None
+        if reload:
+            self.reloadtask = asyncio.ensure_future(self.watcher())
+        else:
+            self.loadscript()
+
+    def done(self):
+        if self.reloadtask:
+            self.reloadtask.cancel()
+
     @property
     def addons(self):
         return [self.ns] if self.ns else []
 
-    def tick(self):
-        if time.time() - self.last_load > self.ReloadInterval:
+    def loadscript(self):
+        ctx.log.info("Loading script %s" % self.path)
+        if self.ns:
+            ctx.master.addons.remove(self.ns)
+        self.ns = None
+        with addonmanager.safecall():
+            ns = load_script(self.fullpath)
+            ctx.master.addons.register(ns)
+            self.ns = ns
+        if self.ns:
+            # We're already running, so we have to explicitly register and
+            # configure the addon
+            ctx.master.addons.invoke_addon(self.ns, "running")
+            ctx.master.addons.invoke_addon(
+                self.ns,
+                "configure",
+                ctx.options.keys()
+            )
+
+    async def watcher(self):
+        last_mtime = 0
+        while True:
             try:
                 mtime = os.stat(self.fullpath).st_mtime
             except FileNotFoundError:
+                ctx.log.info("Removing script %s" % self.path)
                 scripts = list(ctx.options.scripts)
                 scripts.remove(self.path)
                 ctx.options.update(scripts=scripts)
                 return
-
-            if mtime > self.last_mtime:
-                ctx.log.info("Loading script: %s" % self.path)
-                if self.ns:
-                    ctx.master.addons.remove(self.ns)
-                self.ns = None
-                with addonmanager.safecall():
-                    ns = load_script(self.fullpath)
-                    ctx.master.addons.register(ns)
-                    self.ns = ns
-                if self.ns:
-                    # We're already running, so we have to explicitly register and
-                    # configure the addon
-                    ctx.master.addons.invoke_addon(self.ns, "running")
-                    ctx.master.addons.invoke_addon(
-                        self.ns,
-                        "configure",
-                        ctx.options.keys()
-                    )
-                self.last_load = time.time()
-                self.last_mtime = mtime
+            if mtime > last_mtime:
+                self.loadscript()
+                last_mtime = mtime
+            await asyncio.sleep(ReloadInterval)
 
 
 class ScriptLoader:
@@ -125,9 +139,7 @@ class ScriptLoader:
     def load(self, loader):
         loader.add_option(
             "scripts", typing.Sequence[str], [],
-            """
-            Execute a script.
-            """
+            "Execute a script."
         )
 
     def running(self):
@@ -141,12 +153,7 @@ class ScriptLoader:
             simulated.
         """
         try:
-            s = Script(path)
-            l = addonmanager.Loader(ctx.master)
-            ctx.master.addons.invoke_addon(s, "load", l)
-            ctx.master.addons.invoke_addon(s, "configure", ctx.options.keys())
-            # Script is loaded on the first tick
-            ctx.master.addons.invoke_addon(s, "tick")
+            s = Script(path, False)
             for f in flows:
                 for evt, arg in eventsequence.iterate(f):
                     ctx.master.addons.invoke_addon(s, evt, arg)
@@ -161,7 +168,7 @@ class ScriptLoader:
 
             for a in self.addons[:]:
                 if a.path not in ctx.options.scripts:
-                    ctx.log.info("Un-loading script: %s" % a.name)
+                    ctx.log.info("Un-loading script: %s" % a.path)
                     ctx.master.addons.remove(a)
                     self.addons.remove(a)
 
@@ -181,7 +188,7 @@ class ScriptLoader:
                 if s in current:
                     ordered.append(current[s])
                 else:
-                    sc = Script(s)
+                    sc = Script(s, True)
                     ordered.append(sc)
                     newscripts.append(sc)
 
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index bbbd07d0..8eb01600 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -94,19 +94,14 @@ class Master:
         exc = None
         try:
             loop()
-        except Exception as e:
+        except Exception as e:  # pragma: no cover
             exc = traceback.format_exc()
         finally:
-            if not self.should_exit.is_set():
+            if not self.should_exit.is_set():  # pragma: no cover
                 self.shutdown()
-            pending = asyncio.Task.all_tasks()
             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
+            for p in asyncio.Task.all_tasks():
+                p.cancel()
             loop.close()
 
         if exc:  # pragma: no cover
@@ -122,6 +117,7 @@ class Master:
         self.run_loop(loop.run_forever)
 
     async def _shutdown(self):
+        self.should_exit.set()
         if self.server:
             self.server.shutdown()
         loop = asyncio.get_event_loop()
diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py
index 73e57456..e0e4fef9 100644
--- a/mitmproxy/test/taddons.py
+++ b/mitmproxy/test/taddons.py
@@ -123,11 +123,7 @@ class context:
         """
             Loads a script from path, and returns the enclosed addon.
         """
-        sc = script.Script(path)
-        loader = addonmanager.Loader(self.master)
-        self.master.addons.invoke_addon(sc, "load", loader)
-        self.configure(sc)
-        self.master.addons.invoke_addon(sc, "tick")
+        sc = script.Script(path, False)
         return sc.addons[0] if sc.addons else None
 
     def invoke(self, addon, event, *args, **kwargs):
diff --git a/test/mitmproxy/addons/test_onboarding.py b/test/mitmproxy/addons/test_onboarding.py
index 0d99b1ff..a942062f 100644
--- a/test/mitmproxy/addons/test_onboarding.py
+++ b/test/mitmproxy/addons/test_onboarding.py
@@ -4,23 +4,21 @@ from mitmproxy.addons import onboarding
 from mitmproxy.test import taddons
 from .. import tservers
 
-import asyncio
-import tornado.platform.asyncio
-asyncio.set_event_loop_policy(tornado.platform.asyncio.AnyThreadEventLoopPolicy())
-
 
 class TestApp(tservers.HTTPProxyTest):
     def addons(self):
         return [onboarding.Onboarding()]
 
-    def test_basic(self):
+    @pytest.mark.asyncio
+    async def test_basic(self):
         ob = onboarding.Onboarding()
         with taddons.context(ob) as tctx:
             tctx.configure(ob)
             assert self.app("/").status_code == 200
 
     @pytest.mark.parametrize("ext", ["pem", "p12"])
-    def test_cert(self, ext):
+    @pytest.mark.asyncio
+    async def test_cert(self, ext):
         ob = onboarding.Onboarding()
         with taddons.context(ob) as tctx:
             tctx.configure(ob)
@@ -29,7 +27,8 @@ class TestApp(tservers.HTTPProxyTest):
             assert resp.content
 
     @pytest.mark.parametrize("ext", ["pem", "p12"])
-    def test_head(self, ext):
+    @pytest.mark.asyncio
+    async def test_head(self, ext):
         ob = onboarding.Onboarding()
         with taddons.context(ob) as tctx:
             tctx.configure(ob)
diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py
index 96e19841..fad514ae 100644
--- a/test/mitmproxy/addons/test_script.py
+++ b/test/mitmproxy/addons/test_script.py
@@ -12,6 +12,10 @@ from mitmproxy.test import tflow
 from mitmproxy.test import tutils
 
 
+# We want this to be speedy for testing
+script.ReloadInterval = 0.1
+
+
 @pytest.mark.asyncio
 async def test_load_script():
     with taddons.context() as tctx:
@@ -71,7 +75,7 @@ class TestScript:
     def test_notfound(self):
         with taddons.context():
             with pytest.raises(exceptions.OptionsError):
-                script.Script("nonexistent")
+                script.Script("nonexistent", False)
 
     def test_quotes_around_filename(self):
         """
@@ -81,21 +85,23 @@ class TestScript:
         path = tutils.test_data.path("mitmproxy/data/addonscripts/recorder/recorder.py")
 
         s = script.Script(
-            '"{}"'.format(path)
+            '"{}"'.format(path),
+            False
         )
         assert '"' not in s.fullpath
 
-    def test_simple(self):
+    @pytest.mark.asyncio
+    async def test_simple(self):
         with taddons.context() as tctx:
             sc = script.Script(
                 tutils.test_data.path(
                     "mitmproxy/data/addonscripts/recorder/recorder.py"
-                )
+                ),
+                True,
             )
             tctx.master.addons.add(sc)
             tctx.configure(sc)
-            sc.tick()
-
+            await tctx.master.await_log("recorder running")
             rec = tctx.master.addons.get("recorder")
 
             assert rec.call_log[0][0:2] == ("recorder", "load")
@@ -112,25 +118,24 @@ class TestScript:
             f = tmpdir.join("foo.py")
             f.ensure(file=True)
             f.write("\n")
-            sc = script.Script(str(f))
+            sc = script.Script(str(f), True)
             tctx.configure(sc)
-            sc.tick()
             assert await tctx.master.await_log("Loading")
-            tctx.master.clear()
 
-            sc.last_load, sc.last_mtime = 0, 0
-            sc.tick()
+            tctx.master.clear()
+            f.write("\n")
             assert await tctx.master.await_log("Loading")
 
     @pytest.mark.asyncio
     async def test_exception(self):
         with taddons.context() as tctx:
             sc = script.Script(
-                tutils.test_data.path("mitmproxy/data/addonscripts/error.py")
+                tutils.test_data.path("mitmproxy/data/addonscripts/error.py"),
+                True,
             )
             tctx.master.addons.add(sc)
+            await tctx.master.await_log("error running")
             tctx.configure(sc)
-            sc.tick()
 
             f = tflow.tflow(resp=True)
             tctx.master.addons.trigger("request", f)
@@ -138,16 +143,17 @@ class TestScript:
             assert await tctx.master.await_log("ValueError: Error!")
             assert await tctx.master.await_log("error.py")
 
-    def test_addon(self):
+    @pytest.mark.asyncio
+    async def test_addon(self):
         with taddons.context() as tctx:
             sc = script.Script(
                 tutils.test_data.path(
                     "mitmproxy/data/addonscripts/addon.py"
-                )
+                ),
+                True
             )
             tctx.master.addons.add(sc)
-            tctx.configure(sc)
-            sc.tick()
+            await tctx.master.await_log("addon running")
             assert sc.ns.event_log == [
                 'scriptload', 'addonload', 'scriptconfigure', 'addonconfigure'
             ]
@@ -184,7 +190,6 @@ class TestScriptLoader:
             debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
             assert debug == [
                 'recorder load', 'recorder running', 'recorder configure',
-                'recorder tick',
                 'recorder requestheaders', 'recorder request',
                 'recorder responseheaders', 'recorder response'
             ]
@@ -224,17 +229,21 @@ class TestScriptLoader:
                     scripts = ["one", "one"]
                 )
 
-    def test_script_deletion(self):
+    @pytest.mark.asyncio
+    async def test_script_deletion(self):
         tdir = tutils.test_data.path("mitmproxy/data/addonscripts/")
         with open(tdir + "/dummy.py", 'w') as f:
             f.write("\n")
+
         with taddons.context() as tctx:
             sl = script.ScriptLoader()
             tctx.master.addons.add(sl)
             tctx.configure(sl, scripts=[tutils.test_data.path("mitmproxy/data/addonscripts/dummy.py")])
+            await tctx.master.await_log("Loading")
 
             os.remove(tutils.test_data.path("mitmproxy/data/addonscripts/dummy.py"))
-            tctx.invoke(sl, "tick")
+
+            await tctx.master.await_log("Removing")
             assert not tctx.options.scripts
             assert not sl.addons
 
@@ -286,17 +295,14 @@ class TestScriptLoader:
                 'a load',
                 'a running',
                 'a configure',
-                'a tick',
 
                 'b load',
                 'b running',
                 'b configure',
-                'b tick',
 
                 'c load',
                 'c running',
                 'c configure',
-                'c tick',
             ]
 
             tctx.master.clear()
@@ -317,7 +323,7 @@ class TestScriptLoader:
                 'b configure',
             ]
 
-            tctx.master.logs = []
+            tctx.master.clear()
             tctx.configure(
                 sc,
                 scripts = [
@@ -325,9 +331,7 @@ class TestScriptLoader:
                     "%s/a.py" % rec,
                 ]
             )
-            tctx.master.addons.invoke_addon(sc, "tick")
-            await tctx.master.await_log("a tick")
-
+            await tctx.master.await_log("Loading")
             debug = [i.msg for i in tctx.master.logs if i.level == "debug"]
             assert debug == [
                 'c done',
@@ -336,6 +340,4 @@ class TestScriptLoader:
                 'e load',
                 'e running',
                 'e configure',
-                'e tick',
-                'a tick',
             ]
diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py
index 8c834d82..c6b540d4 100644
--- a/test/mitmproxy/data/addonscripts/addon.py
+++ b/test/mitmproxy/data/addonscripts/addon.py
@@ -1,3 +1,4 @@
+from mitmproxy import ctx
 event_log = []
 
 
@@ -7,6 +8,7 @@ class Addon:
         return event_log
 
     def load(self, opts):
+        ctx.log.info("addon running")
         event_log.append("addonload")
 
     def configure(self, updated):
diff --git a/test/mitmproxy/data/addonscripts/error.py b/test/mitmproxy/data/addonscripts/error.py
index 4a3c370f..2f0c1755 100644
--- a/test/mitmproxy/data/addonscripts/error.py
+++ b/test/mitmproxy/data/addonscripts/error.py
@@ -1,6 +1,9 @@
-def mkerr():
-    raise ValueError("Error!")
+from mitmproxy import ctx
+
+
+def running():
+    ctx.log.info("error running")
 
 
 def request(flow):
-    mkerr()
+    raise ValueError("Error!")
diff --git a/test/mitmproxy/data/addonscripts/shutdown.py b/test/mitmproxy/data/addonscripts/shutdown.py
new file mode 100644
index 00000000..f4a8f55d
--- /dev/null
+++ b/test/mitmproxy/data/addonscripts/shutdown.py
@@ -0,0 +1,5 @@
+from mitmproxy import ctx
+
+
+def running():
+    ctx.master.shutdown()
\ No newline at end of file
diff --git a/test/mitmproxy/data/addonscripts/stream_modify.py b/test/mitmproxy/data/addonscripts/stream_modify.py
index 4fbf45c2..4ebcc3e9 100644
--- a/test/mitmproxy/data/addonscripts/stream_modify.py
+++ b/test/mitmproxy/data/addonscripts/stream_modify.py
@@ -1,7 +1,13 @@
+from mitmproxy import ctx
+
 def modify(chunks):
     for chunk in chunks:
         yield chunk.replace(b"foo", b"bar")
 
 
+def running():
+    ctx.log.info("stream_modify running")
+
+
 def responseheaders(flow):
     flow.response.stream = modify
diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py
index bf24e28b..d7836104 100644
--- a/test/mitmproxy/proxy/test_server.py
+++ b/test/mitmproxy/proxy/test_server.py
@@ -256,11 +256,15 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin):
             resp = p.request("get:'http://foo':h':foo'='bar'")
         assert resp.status_code == 400
 
-    def test_stream_modify(self):
+    @pytest.mark.asyncio
+    async def test_stream_modify(self):
         s = script.Script(
-            tutils.test_data.path("mitmproxy/data/addonscripts/stream_modify.py")
+            tutils.test_data.path("mitmproxy/data/addonscripts/stream_modify.py"),
+            False,
         )
         self.set_addons(s)
+        await self.master.await_log("stream_modify running")
+
         d = self.pathod('200:b"foo"')
         assert d.content == b"bar"
 
@@ -564,7 +568,8 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin):
 
     def test_tcp_stream_modify(self):
         s = script.Script(
-            tutils.test_data.path("mitmproxy/data/addonscripts/tcp_stream_modify.py")
+            tutils.test_data.path("mitmproxy/data/addonscripts/tcp_stream_modify.py"),
+            False,
         )
         self.set_addons(s)
         self._tcpproxy_on()
diff --git a/test/mitmproxy/tools/test_main.py b/test/mitmproxy/tools/test_main.py
index 57544276..751dcbe0 100644
--- a/test/mitmproxy/tools/test_main.py
+++ b/test/mitmproxy/tools/test_main.py
@@ -1,19 +1,23 @@
-import pytest
+import asyncio
 
 from mitmproxy.tools import main
-from mitmproxy import ctx
+from mitmproxy.test import tutils
 
+shutdown_script = tutils.test_data.path("mitmproxy/data/addonscripts/shutdown.py")
 
-@pytest.mark.asyncio
-async def test_mitmweb(event_loop):
+
+def test_mitmweb(event_loop):
+    asyncio.set_event_loop(event_loop)
     main.mitmweb([
         "--no-web-open-browser",
+        "-s", shutdown_script,
         "-q", "-p", "0",
     ])
-    await ctx.master._shutdown()
 
 
-@pytest.mark.asyncio
-async def test_mitmdump():
-    main.mitmdump(["-q", "-p", "0"])
-    await ctx.master._shutdown()
+def test_mitmdump(event_loop):
+    asyncio.set_event_loop(event_loop)
+    main.mitmdump([
+        "-s", shutdown_script,
+        "-q", "-p", "0",
+    ])
diff --git a/test/mitmproxy/tools/web/test_static_viewer.py b/test/mitmproxy/tools/web/test_static_viewer.py
index dfc45bc2..c044dee8 100644
--- a/test/mitmproxy/tools/web/test_static_viewer.py
+++ b/test/mitmproxy/tools/web/test_static_viewer.py
@@ -1,5 +1,6 @@
 import json
 from unittest import mock
+import pytest
 
 from mitmproxy.test import taddons
 from mitmproxy.test import tflow
@@ -57,7 +58,8 @@ def test_save_flows_content(ctx, tmpdir):
         assert p.join('response/content/Auto.json').check(file=1)
 
 
-def test_static_viewer(tmpdir):
+@pytest.mark.asyncio
+async def test_static_viewer(tmpdir):
     s = static_viewer.StaticViewer()
     rf = readfile.ReadFile()
     sa = save.Save()