Improve benchmarking

- The benchmark addon now manages setting up and tearing down the backend and
traffic processes itself.
- Use wrk instead of hey. I get more consistent results with this tool, and hey
shows a strange tail-latency bump that seems artificial.
- Make termination behaviour simpler. The bencmark revealed a bug where .done
events were not called if the proxy was shut down by an addon.

7 files changed, 64 insertions, 75 deletions

diff --git a/test/bench/README.md b/test/bench/README.md
index 05741c07..3d9e7ef7 100644
--- a/test/bench/README.md
+++ b/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.
 
diff --git a/test/bench/backend b/test/bench/backend
deleted file mode 100755
index 12a05d70..00000000
--- a/test/bench/backend
+++ /dev/null
@@ -1,3 +0,0 @@
-#!/bin/sh
-
-devd -p 10001 .
\ No newline at end of file
diff --git a/test/bench/benchmark.py b/test/bench/benchmark.py
new file mode 100644
index 00000000..8d208088
--- /dev/null
+++ b/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()]
\ No newline at end of file
diff --git a/test/bench/profiler.py b/test/bench/profiler.py
deleted file mode 100644
index 9072e17d..00000000
--- a/test/bench/profiler.py
+++ /dev/null
@@ -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()]
\ No newline at end of file
diff --git a/test/bench/run b/test/bench/run
new file mode 100755
index 00000000..9925b578
--- /dev/null
+++ b/test/bench/run
@@ -0,0 +1,4 @@
+#!/bin/sh
+
+mkdir -p results
+mitmdump -p0 -q --set benchmark_save_path=./results/mitmdump -s ./benchmark.py
\ No newline at end of file
diff --git a/test/bench/simple.mitmproxy b/test/bench/simple.mitmproxy
deleted file mode 100755
index 9de32981..00000000
--- a/test/bench/simple.mitmproxy
+++ /dev/null
@@ -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
diff --git a/test/bench/simple.traffic b/test/bench/simple.traffic
deleted file mode 100755
index 08200e05..00000000
--- a/test/bench/simple.traffic
+++ /dev/null
@@ -1,3 +0,0 @@
-#!/bin/sh
-
-hey -disable-keepalive http://localhost:10002/profiler.py | tee ./results/simple.perf
\ No newline at end of file