Merge pull request #3048 from cortesi/qbench

Improve benchmarking

mitmproxy/controller.py

@@ -21,26 +21,28 @@ class Channel:
         Raises:
             exceptions.Kill: All connections should be closed immediately.
         """
-        m.reply = Reply(m)
-        asyncio.run_coroutine_threadsafe(
-            self.master.addons.handle_lifecycle(mtype, m),
-            self.loop,
-        )
-        g = m.reply.q.get()
-        if g == exceptions.Kill:
-            raise exceptions.Kill()
-        return g
+        if not self.should_exit.is_set():
+            m.reply = Reply(m)
+            asyncio.run_coroutine_threadsafe(
+                self.master.addons.handle_lifecycle(mtype, m),
+                self.loop,
+            )
+            g = m.reply.q.get()
+            if g == exceptions.Kill:
+                raise exceptions.Kill()
+            return g
 
     def tell(self, mtype, m):
         """
         Decorate a message with a dummy reply attribute, send it to the master,
         then return immediately.
         """
-        m.reply = DummyReply()
-        asyncio.run_coroutine_threadsafe(
-            self.master.addons.handle_lifecycle(mtype, m),
-            self.loop,
-        )
+        if not self.should_exit.is_set():
+            m.reply = DummyReply()
+            asyncio.run_coroutine_threadsafe(
+                self.master.addons.handle_lifecycle(mtype, m),
+                self.loop,
+            )
 
 
 NO_REPLY = object()  # special object we can distinguish from a valid "None" reply.

mitmproxy/master.py

@@ -81,7 +81,6 @@ class Master:
             self.addons.trigger("running")
         while True:
             if self.should_exit.is_set():
-                asyncio.get_event_loop().stop()
                 return
             self.addons.trigger("tick")
             await asyncio.sleep(0.1)
@@ -93,14 +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")
 
-    def shutdown(self):
+    async def _shutdown(self):
         if self.server:
             self.server.shutdown()
-        self.should_exit.set()
+        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():
+            self.should_exit.set()
+            asyncio.run_coroutine_threadsafe(
+                self._shutdown(),
+                loop = self.channel.loop,
+            )
 
     def _change_reverse_host(self, f):
         """

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)

test/bench/README.md

@@ -1,7 +1,7 @@
 
-This directory contains a set of tools for benchmarking and profiling mitmproxy.
-At the moment, this is simply to give developers a quick way to see the impact
-of their work. Eventually, this might grow into a performance dashboard with
+This directory contains an addon for benchmarking and profiling mitmproxy. At
+the moment, this is simply to give developers a quick way to see the impact of
+their work. Eventually, this might grow into a performance dashboard with
 historical data, so we can track performance over time.
 
 
@@ -9,48 +9,18 @@ historical data, so we can track performance over time.
 
 Install the following tools:
 
-    go get -u github.com/rakyll/hey
+    https://github.com/wg/wrk
+
     go get github.com/cortesi/devd/cmd/devd
 
 You may also want to install snakeviz to make viewing profiles easier:
 
     pip install snakeviz
 
-In one window, run the devd server:
-
-    ./backend
-
-
-# Running tests
-
-Each run consists of two files - a mitproxy invocation, and a traffic generator.
-Make sure the backend is started, then run the proxy:
-
-    ./simple.mitmproxy
-
-Now run the traffic generator:
-
-    ./simple.traffic
-
-After the run is done, quit the proxy with ctrl-c.
-
-
-# Reading results
-
-Results are placed in the ./results directory. You should see two files - a
-performance log from **hey**, and a profile. You can view the profile like so:
-
-    snakeviz ./results/simple.prof
-
-
-
-
-
-
-
-
-
-
+Now run the benchmark by loading the addon. A typical invocation is as follows:
 
+    mitmdump -p0 -q --set benchmark_save_path=/tmp/foo -s ./benchmark.py
 
+This will start up the backend server, run the benchmark, save the results to
+/tmp/foo.bench and /tmp/foo.prof, and exit.
 

test/bench/backend

@@ -1,3 +0,0 @@
-#!/bin/sh
-
-devd -p 10001 .

test/bench/benchmark.py

@@ -0,0 +1,51 @@
+import asyncio
+import cProfile
+from mitmproxy import ctx
+
+
+class Benchmark:
+    """
+        A simple profiler addon.
+    """
+    def __init__(self):
+        self.pr = cProfile.Profile()
+        self.started = False
+
+    async def procs(self):
+        ctx.log.error("starting benchmark")
+        backend = await asyncio.create_subprocess_exec("devd", "-q", "-p", "10001", ".")
+        traf = await asyncio.create_subprocess_exec(
+            "wrk",
+            "-c50",
+            "-d5s",
+            "http://localhost:%s/benchmark.py" % ctx.master.server.address[1],
+            stdout=asyncio.subprocess.PIPE
+        )
+        stdout, _ = await traf.communicate()
+        open(ctx.options.benchmark_save_path + ".bench", mode="wb").write(stdout)
+        ctx.log.error(stdout.decode("ascii"))
+        backend.kill()
+        ctx.master.shutdown()
+
+    def load(self, loader):
+        loader.add_option(
+            "benchmark_save_path",
+            str,
+            "/tmp/profile",
+            "Destination for the .prof and and .bench result files"
+        )
+        ctx.options.update(
+            mode="reverse:http://devd.io:10001",
+        )
+        self.pr.enable()
+
+    def running(self):
+        if not self.started:
+            self.started = True
+            asyncio.get_event_loop().create_task(self.procs())
+
+    def done(self):
+        self.pr.dump_stats(ctx.options.benchmark_save_path + ".prof")
+
+
+addons = [Benchmark()]

test/bench/profiler.py

@@ -1,25 +0,0 @@
-import cProfile
-from mitmproxy import ctx
-
-
-class Profile:
-    """
-        A simple profiler addon.
-    """
-    def __init__(self):
-        self.pr = cProfile.Profile()
-
-    def load(self, loader):
-        loader.add_option(
-            "profile_path",
-            str,
-            "/tmp/profile",
-            "Destination for the run profile, saved at exit"
-        )
-        self.pr.enable()
-
-    def done(self):
-        self.pr.dump_stats(ctx.options.profile_path)
-
-
-addons = [Profile()]

test/bench/run

@@ -0,0 +1,4 @@
+#!/bin/sh
+
+mkdir -p results
+mitmdump -p0 -q --set benchmark_save_path=./results/mitmdump -s ./benchmark.py

test/bench/simple.mitmproxy

@@ -1,5 +0,0 @@
-#!/bin/sh
-
-mkdir -p results
-mitmdump -p 10002 --mode reverse:http://devd.io:10001 \
-    -s ./profiler.py --set profile_path=./results/simple.prof

test/bench/simple.traffic

@@ -1,3 +0,0 @@
-#!/bin/sh
-
-hey -disable-keepalive http://localhost:10002/profiler.py | tee ./results/simple.perf

test/helper_tools/bench.py

@@ -1,23 +0,0 @@
-import requests
-import time
-
-n = 100
-url = "http://192.168.1.1/"
-proxy = "http://192.168.1.115:8080/"
-
-start = time.time()
-for _ in range(n):
-    requests.get(url, allow_redirects=False, proxies=dict(http=proxy))
-    print(".", end="")
-t_mitmproxy = time.time() - start
-
-print("\r\nTotal time with mitmproxy: {}".format(t_mitmproxy))
-
-
-start = time.time()
-for _ in range(n):
-    requests.get(url, allow_redirects=False)
-    print(".", end="")
-t_without = time.time() - start
-
-print("\r\nTotal time without mitmproxy: {}".format(t_without))

test/helper_tools/benchtool.py

@@ -1,56 +0,0 @@
-# Profile mitmdump with apachebench and
-# yappi (https://code.google.com/p/yappi/)
-#
-# Requirements:
-# - Apache Bench "ab" binary
-# - pip install click yappi
-
-from mitmproxy.main import mitmdump
-from os import system
-from threading import Thread
-import time
-
-import yappi
-import click
-
-
-class ApacheBenchThread(Thread):
-
-    def __init__(self, concurrency):
-        self.concurrency = concurrency
-        super().__init__()
-
-    def run(self):
-        time.sleep(2)
-        system(
-            "ab -n 1024 -c {} -X 127.0.0.1:8080 http://example.com/".format(self.concurrency))
-
-
-@click.command()
-@click.option('--profiler', default="none", type=click.Choice(['none', 'yappi']))
-@click.option('--clock-type', default="cpu", type=click.Choice(['wall', 'cpu']))
-@click.option('--concurrency', default=1, type=click.INT)
-def main(profiler, clock_type, concurrency):
-
-    outfile = "callgrind.mitmdump-{}-c{}".format(clock_type, concurrency)
-    a = ApacheBenchThread(concurrency)
-    a.start()
-
-    if profiler == "yappi":
-        yappi.set_clock_type(clock_type)
-        yappi.start(addons=True)
-
-    print("Start mitmdump...")
-    mitmdump(["-k", "-q", "-S", "1024example"])
-    print("mitmdump stopped.")
-
-    print("Save profile information...")
-    if profiler == "yappi":
-        yappi.stop()
-        stats = yappi.get_func_stats()
-        stats.save(outfile, type='callgrind')
-    print("Done.")
-
-
-if __name__ == '__main__':
-    main()

test/helper_tools/testpatt

@@ -1,9 +0,0 @@
-#!/bin/bash
-
-# Generate a test pattern with pathoc
-PATHOD=localhost:9999
-pathoc -s -c $PATHOD localhost:8080 "get:'/p/200:p0,1:b@2048b':b@2048b"
-pathoc -s -c $PATHOD localhost:8080 "get:'/p/300:p0,1:b@2048b':b@2048b"
-pathoc -s -c $PATHOD localhost:8080 "get:'/p/400:p0,1:b@2048b':b@2048b"
-pathoc -s -c $PATHOD localhost:8080 "get:'/p/500:p0,1:b@2048b':b@2048b"
-pathoc -s -c $PATHOD localhost:8080 "get:'/p/600:p0,1:b@2048b':b@2048b"

test/helper_tools/typehints_for_options.py

@@ -1,34 +0,0 @@
-import typing
-from unittest import mock
-
-from mitmproxy import proxy, options
-from mitmproxy.tools import dump, console, web
-
-
-def print_typehints(opts):
-    for name, option in sorted(opts.items()):
-        print(
-            # For Python 3.6, we can just use "{}: {}".
-            "{}: {} = None".format(
-                name,
-                {
-                    int: "int",
-                    str: "str",
-                    bool: "bool",
-                    typing.Optional[str]: "Optional[str]",
-                    typing.Sequence[str]: "Sequence[str]"
-                }[option.typespec]
-            )
-        )
-
-
-if __name__ == "__main__":
-    opts = options.Options()
-    server = proxy.server.DummyServer(None)
-
-    # initialize with all three tools here to capture tool-specific options defined in addons.
-    dump.DumpMaster(opts, server)
-    with mock.patch("sys.stdout.isatty", lambda: True):
-        console.master.ConsoleMaster(opts, server)
-    web.master.WebMaster(opts, server)
-    print_typehints(opts)

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

test/mitmproxy/data/addonscripts/shutdown.py

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

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

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: