asyncio: ditch the handler context

There are a few reasons for this. First, logs are now async, and can be called
at any time. Second, the event loop is thread local, so there can only ever be
one master per thread. These two things together completely obviate the need
for a handler context.

7 files changed, 21 insertions, 54 deletions

diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py
index 7c446c80..9d0e7069 100644
--- a/mitmproxy/addonmanager.py
+++ b/mitmproxy/addonmanager.py
@@ -8,7 +8,6 @@ from mitmproxy import exceptions
 from mitmproxy import eventsequence
 from mitmproxy import controller
 from mitmproxy import flow
-from mitmproxy import log
 from . import ctx
 import pprint
 
@@ -38,9 +37,6 @@ def cut_traceback(tb, func_name):
 
 
 class StreamLog:
-    """
-        A class for redirecting output using contextlib.
-    """
     def __init__(self, lg):
         self.log = lg
 
@@ -183,9 +179,8 @@ class AddonManager:
             Add addons to the end of the chain, and run their load event.
             If any addon has sub-addons, they are registered.
         """
-        with self.master.handlecontext():
-            for i in addons:
-                self.chain.append(self.register(i))
+        for i in addons:
+            self.chain.append(self.register(i))
 
     def remove(self, addon):
         """
@@ -201,8 +196,7 @@ class AddonManager:
                 raise exceptions.AddonManagerError("No such addon: %s" % n)
             self.chain = [i for i in self.chain if i is not a]
             del self.lookup[_get_name(a)]
-        with self.master.handlecontext():
-            self.invoke_addon(a, "done")
+        self.invoke_addon(a, "done")
 
     def __len__(self):
         return len(self.chain)
@@ -245,8 +239,7 @@ class AddonManager:
 
     def invoke_addon(self, addon, name, *args, **kwargs):
         """
-            Invoke an event on an addon and all its children. This method must
-            run within an established handler context.
+            Invoke an event on an addon and all its children.
         """
         if name not in eventsequence.Events:
             name = "event_" + name
@@ -268,12 +261,11 @@ class AddonManager:
 
     def trigger(self, name, *args, **kwargs):
         """
-            Establish a handler context and trigger an event across all addons
+            Trigger an event across all addons.
         """
-        with self.master.handlecontext():
-            for i in self.chain:
-                try:
-                    with safecall():
-                        self.invoke_addon(i, name, *args, **kwargs)
-                except exceptions.AddonHalt:
-                    return
+        for i in self.chain:
+            try:
+                with safecall():
+                    self.invoke_addon(i, name, *args, **kwargs)
+            except exceptions.AddonHalt:
+                return
diff --git a/mitmproxy/command.py b/mitmproxy/command.py
index 114e882d..35a7f42a 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -95,11 +95,7 @@ class Command:
             Call the command with a list of arguments. At this point, all
             arguments are strings.
         """
-        pargs = self.prepare_args(args)
-
-        with self.manager.master.handlecontext():
-            ret = self.func(*pargs)
-
+        ret = self.func(*self.prepare_args(args))
         if ret is None and self.returntype is None:
             return
         typ = mitmproxy.types.CommandTypes.get(self.returntype)
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 323f73dc..4e17e538 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -1,5 +1,4 @@
 import threading
-import contextlib
 import asyncio
 import logging
 
@@ -58,6 +57,10 @@ class Master:
         self.waiting_flows = []
         self.log = log.Log(self)
 
+        mitmproxy_ctx.master = self
+        mitmproxy_ctx.log = self.log
+        mitmproxy_ctx.options = self.options
+
     @property
     def server(self):
         return self._server
@@ -67,22 +70,6 @@ class Master:
         server.set_channel(self.channel)
         self._server = server
 
-    @contextlib.contextmanager
-    def handlecontext(self):
-        # Handlecontexts also have to nest - leave cleanup to the outermost
-        if mitmproxy_ctx.master:
-            yield
-            return
-        mitmproxy_ctx.master = self
-        mitmproxy_ctx.log = self.log
-        mitmproxy_ctx.options = self.options
-        try:
-            yield
-        finally:
-            mitmproxy_ctx.master = None
-            mitmproxy_ctx.log = None
-            mitmproxy_ctx.options = None
-
     def start(self):
         self.should_exit.clear()
         if self.server:
diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py
index 6d439c4a..73e57456 100644
--- a/mitmproxy/test/taddons.py
+++ b/mitmproxy/test/taddons.py
@@ -74,7 +74,6 @@ class context:
             options
         )
         self.options = self.master.options
-        self.wrapped = None
 
         if loadcore:
             self.master.addons.add(core.Core())
@@ -82,20 +81,10 @@ class context:
         for a in addons:
             self.master.addons.add(a)
 
-    def ctx(self):
-        """
-            Returns a new handler context.
-        """
-        return self.master.handlecontext()
-
     def __enter__(self):
-        self.wrapped = self.ctx()
-        self.wrapped.__enter__()
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        self.wrapped.__exit__(exc_type, exc_value, traceback)
-        self.wrapped = None
         return False
 
     @contextlib.contextmanager
diff --git a/setup.py b/setup.py
index 18fb3827..3e914847 100644
--- a/setup.py
+++ b/setup.py
@@ -88,6 +88,7 @@ setup(
             "flake8>=3.5, <3.6",
             "Flask>=0.10.1, <0.13",
             "mypy>=0.580,<0.581",
+            "pytest-asyncio>=0.8",
             "pytest-cov>=2.5.1,<3",
             "pytest-faulthandler>=1.3.1,<2",
             "pytest-timeout>=1.2.1,<2",
diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py
index 17b1e372..3f0ce68c 100644
--- a/test/mitmproxy/addons/test_script.py
+++ b/test/mitmproxy/addons/test_script.py
@@ -180,11 +180,12 @@ class TestScriptLoader:
                 'recorder responseheaders', 'recorder response'
             ]
 
-    def test_script_run_nonexistent(self):
+    @pytest.mark.asyncio
+    async def test_script_run_nonexistent(self):
         sc = script.ScriptLoader()
         with taddons.context(sc) as tctx:
             sc.script_run([tflow.tflow(resp=True)], "/")
-            tctx.master.has_log("/: No such script")
+            assert await tctx.master.await_log("/: No such script")
 
     def test_simple(self):
         sc = script.ScriptLoader()
diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py
index 347c0dfc..3ce1436a 100644
--- a/test/mitmproxy/proxy/protocol/test_websocket.py
+++ b/test/mitmproxy/proxy/protocol/test_websocket.py
@@ -47,6 +47,7 @@ class _WebSocketServerBase(net_tservers.ServerTestBase):
 
 
 class _WebSocketTestBase:
+    client = None
 
     @classmethod
     def setup_class(cls):