diff options
| -rw-r--r-- | mitmproxy/controller.py | 30 | ||||
| -rw-r--r-- | mitmproxy/master.py | 12 | ||||
| -rw-r--r-- | test/bench/README.md | 48 | ||||
| -rwxr-xr-x | test/bench/backend | 3 | ||||
| -rw-r--r-- | test/bench/benchmark.py | 51 | ||||
| -rw-r--r-- | test/bench/profiler.py | 25 | ||||
| -rwxr-xr-x | test/bench/run | 4 | ||||
| -rwxr-xr-x | test/bench/simple.mitmproxy | 5 | ||||
| -rwxr-xr-x | test/bench/simple.traffic | 3 |
9 files changed, 88 insertions, 93 deletions
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 3b4c3092..582a6683 100644 --- a/mitmproxy/controller.py +++ b/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. diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 51aed1e0..dff8ca5c 100644 --- a/mitmproxy/master.py +++ b/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) @@ -94,13 +93,18 @@ class Master: 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): - if self.server: - self.server.shutdown() - self.should_exit.set() + if not self.should_exit.is_set(): + if self.server: + self.server.shutdown() + self.should_exit.set() + loop = asyncio.get_event_loop() + loop.stop() def _change_reverse_host(self, f): """ 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 |