Repository: mitmproxy/mitmproxy

Issue:
Improve flow list performance
mitmproxy's flow list is very slow: holding the down arrow to move the cursor causes it to visibly lag on my machine. The lag affects not just the UI, but the entire application - requests are slowed down as well.

Profiling (using `cProfile`) identified the bottlenecks as:

- For every redraw, `FlowListWalker` was creating several instances of `FlowList` *per flow* for every redraw. This is because `urwid.ListBox` needs to iterate over the list more than once per redraw.

- `FlowListWalker` uses `commands.execute` instead of calling `View` methods directly. The formatting and parsing of commands adds an additional performance cost. (This is a recent performance regression only in git, introduced by @madt1m in commit 773c9535146991af38f954cb4d3b0ead9d7485c5.)

- `raw_format_flow` incorrectly had the entire `Flow` object added to its argument list. This caused needless work for `@lru_cache`, as it was hashing the entire `Flow` object as part of the cache key, and probably causing unnecessary cache misses as `Flow` properties changed without affecting the displayed content. (This is a performance regression, introduced by @Kriechi in commit d86cb76e5ba38c51b7d6017fb151e4946e3fb916.)

Profiling result:

![](https://raw.githubusercontent.com/BkPHcgQL3V/mitmproxy/files/profile.svg?sanitize=true)

The changes in this pull request rectify these problems, thus improving performance somewhat.


Discussion:
Thanks for tracking this down! This looks like amazing work!

I do not recall the reasoning for changing `raw_format_flow` - so your guess is as good as mine.

If you could make our linting system happy as well, I think we can merge:
`mitmproxy/tools/console/flowlist.py:38: error: Need type annotation for '__cache'`
Ehi there! Nice piece of work you did here. Actually, there's a reason behind that change: the general assumption at the time of commit was that directly accessing the View addon from other tools is somewhat hacky, and this eventually lead to a Codebase clean-up, to make tools and add-ons communicate between themselves only through our command layer. 
> This looks like amazing work!

Thanks!

> I do not recall the reasoning for changing `raw_format_flow` - so your guess is as good as mine.

From the commit in question (d86cb76e5ba38c51b7d6017fb151e4946e3fb916), it looks like at the time it wasn't obvious that all information about the flow was intended to be passed through the tuple parameter, so that `@lru_cache` would use it as a cache key.

> the general assumption at the time of commit was that directly accessing the View addon from other tools is somewhat hacky, and this eventually lead to a Codebase clean-up, to make tools and add-ons communicate between themselves only through our command layer.

Yes, that makes sense; unfortunately, this particular change turned out to be on a hot path, so directly accessing internals may be worthwhile.

> Looking a bit further down the profile graph, it seems like we could significantly alleviate this by lru-caching `command.verify_arg_signature` (14.61%) and `command.lexer` (5.3%)? That would also benefit other parts of mitmproxy.

Those sound like good directions for research. I'll need to figure out how to reliably profile this code path, so that I can get quantitative measurements of any improvements.

> How about we just use `functools.lru_cache(maxsize=None)` for `__get` and then call `self.__get.cache_clear()` in `view_changed`?

Sounds good, I'll do this for now.
> I'll need to figure out how to reliably profile this code path, so that I can get quantitative measurements of any improvements.

I wrote the following mitmproxy script:

```python
from mitmproxy import ctx
import timeit

NUM_FLOWS = 100
NUM_SAMPLES = 10
NUM_ITERS = 1


def do_nav():
    for i in range(NUM_FLOWS):
        ctx.master.commands.call("console.nav.down")
    for i in range(NUM_FLOWS):
        ctx.master.commands.call("console.nav.up")


def do_bench(func):
    smallest_time = float('inf')
    for n in range(NUM_SAMPLES):
        time = timeit.timeit(func, number=NUM_ITERS)
        if smallest_time > time:
            smallest_time = time
    return smallest_time


ctx.log.alert('Measuring overhead...')
overhead = do_bench(do_nav)

for i in range(NUM_FLOWS):
    ctx.master.commands.call("view.flows.create", "get", "https://example.com/" + str(i))


ctx.log.alert('Benchmarking...')
time_all = do_bench(do_nav)
time = time_all - overhead


class ResultPrinter:
    def done(self):
        print('Result: ' + str(time) + ' (overhead ' + str(overhead) + ')')


addons = [
    ResultPrinter()
]

ctx.master.commands.call("console.exit")
```

With the above, I now have a clearer picture of which changes are actually worthwhile.

The first commit does not register a measurable performance improvement (I'm guessing I made a mistake when assessing the change), so we can probably drop it. The third commit also does not register an improvement in this benchmark, however the benchmark does not attempt to simulate false cache misses due to things changing in the flow that would not be reflected in `raw_format_flow`'s first argument. Not sure how often those might occur in practice, but I would argue in its favour at least for pedantic correctness' sake.

The second commit, however, does result in a very significant speedup, going from 3.09 to 1.66 seconds according to the above script.

> Looking a bit further down the profile graph, it seems like we could significantly alleviate this by lru-caching `command.verify_arg_signature` (14.61%) and `command.lexer` (5.3%)? That would also benefit other parts of mitmproxy.

It doesn't look like these changes would be as trivial. `command.verify_arg_signature` has a `list` in its argument list, which is apparently not compatible with `@lru_cache`. `command.lexer` merely returns a lexer, so some refactoring would be required to memoise it.
> It doesn't look like these changes would be as trivial. `command.verify_arg_signature` has a `list` in its argument list, which is apparently not compatible with `@lru_cache`. `command.lexer` merely returns a lexer, so some refactoring would be required to memoise it.

`command.lexer`: Yeah, I misread the callgraph. It shouldn't be too difficult though - we'd need to move this chunk into its own function and cache it, which doesn't seem to have external dependencies. https://github.com/mitmproxy/mitmproxy/blob/e2bcca47b1ad8040451cbd95039acf200e9b0e84/mitmproxy/command.py#L160-L176

---------------

For `command.verify_arg_signature`, we can replace
https://github.com/mitmproxy/mitmproxy/blob/e2bcca47b1ad8040451cbd95039acf200e9b0e84/mitmproxy/command.py#L87
with 
```
verify_arg_signature(self.func, tuple(args), None)
```
and
https://github.com/mitmproxy/mitmproxy/blob/e2bcca47b1ad8040451cbd95039acf200e9b0e84/mitmproxy/command.py#L274
with 
```
verify_arg_signature(function, args, kwargs or None)
```

Could you give this a spin on your machine to have a consistent test setup? Thanks so much! 😃 😊 
OK, I'm testing this patch:

```diff
diff --git a/mitmproxy/command.py b/mitmproxy/command.py
index 27f0921d..5283f328 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -14,10 +14,11 @@ from mitmproxy import exceptions
 import mitmproxy.types
 
 
-def verify_arg_signature(f: typing.Callable, args: list, kwargs: dict) -> None:
+@functools.lru_cache(maxsize=200)
+def verify_arg_signature(f: typing.Callable, args: tuple, kwargs: tuple) -> None:
     sig = inspect.signature(f)
     try:
-        sig.bind(*args, **kwargs)
+        sig.bind(*list(args), **dict(kwargs))
     except TypeError as v:
         raise exceptions.CommandError("command argument mismatch: %s" % v.args[0])
 
@@ -31,6 +32,28 @@ def lexer(s):
     return lex
 
 
+@functools.lru_cache(maxsize=200)
+def lex_string(cmdstr):
+    buf = io.StringIO(cmdstr)
+    parts: typing.List[str] = []
+    lex = lexer(buf)
+    while 1:
+        remainder = cmdstr[buf.tell():]
+        try:
+            t = lex.get_token()
+        except ValueError:
+            parts.append(remainder)
+            break
+        if not t:
+            break
+        parts.append(t)
+    if not parts:
+        parts = [""]
+    elif cmdstr.endswith(" "):
+        parts.append("")
+    return parts
+
+
 def typename(t: type) -> str:
     """
         Translates a type to an explanatory string.
@@ -84,7 +107,7 @@ class Command:
         return "%s %s%s" % (self.path, params, ret)
 
     def prepare_args(self, args: typing.Sequence[str]) -> typing.List[typing.Any]:
-        verify_arg_signature(self.func, list(args), {})
+        verify_arg_signature(self.func, tuple(args), ())
 
         remainder: typing.Sequence[str] = []
         if self.has_positional:
@@ -157,23 +180,7 @@ class CommandManager(mitmproxy.types._CommandBase):
         """
             Parse a possibly partial command. Return a sequence of ParseResults and a sequence of remainder type help items.
         """
-        buf = io.StringIO(cmdstr)
-        parts: typing.List[str] = []
-        lex = lexer(buf)
-        while 1:
-            remainder = cmdstr[buf.tell():]
-            try:
-                t = lex.get_token()
-            except ValueError:
-                parts.append(remainder)
-                break
-            if not t:
-                break
-            parts.append(t)
-        if not parts:
-            parts = [""]
-        elif cmdstr.endswith(" "):
-            parts.append("")
+        parts = lex_string(cmdstr)
 
         parse: typing.List[ParseResult] = []
         params: typing.List[type] = []
@@ -271,7 +278,7 @@ def command(path):
     def decorator(function):
         @functools.wraps(function)
         def wrapper(*args, **kwargs):
-            verify_arg_signature(function, args, kwargs)
+            verify_arg_signature(function, args, tuple(kwargs.items()))
             return function(*args, **kwargs)
         wrapper.__dict__["command_path"] = path
         return wrapper
```

Applied to master, it brings the benchmark result from 2.98 to 1.96, which is quite significant! By itself, it's not as great as the second commit in this PR, though. If the patch looks good, I can add it to this PR too.
Awesome, thanks! 🍰 

The diff you posted is uncontroversial I think and I'd be more than happy to merge it right away. Would you mind opening a separate PR for that? 

Your second commit breaks some isolations we wanted to establish so it might be a bit controversial. I have no particular opinion on that tradeoff, but @cortesi might. @cortesi, any thoughts?

I cherry-picked your third commit onto master. :-)
> Would you mind opening a separate PR for that?

OK, here it is: https://github.com/mitmproxy/mitmproxy/pull/3428

> Your second commit breaks some isolations we wanted to establish so it might be a bit controversial. 

Perhaps an additional, abstract interface of sending commands using a more performance-friendly model (tuples and identifiers instead of strings)?

Resolving patch:
diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -1,8 +1,11 @@
+import typing
+from functools import lru_cache
+
 import urwid
 
+import mitmproxy.tools.console.master
 from mitmproxy.tools.console import common
 from mitmproxy.tools.console import layoutwidget
-import mitmproxy.tools.console.master
 
 
 class FlowItem(urwid.WidgetWrap):
@@ -40,6 +43,7 @@ def keypress(self, size, key):
 
 
 class FlowListWalker(urwid.ListWalker):
+    master: "mitmproxy.tools.console.master.ConsoleMaster"
 
     def __init__(self, master):
         self.master = master
@@ -47,13 +51,14 @@ def __init__(self, master):
     def positions(self, reverse=False):
         # The stub implementation of positions can go once this issue is resolved:
         # https://github.com/urwid/urwid/issues/294
-        ret = range(self.master.commands.execute("view.properties.length"))
+        ret = range(self.master.view.get_length())
         if reverse:
             return reversed(ret)
         return ret
 
     def view_changed(self):
         self._modified()
+        self._get.cache_clear()
 
     def get_focus(self):
         if not self.master.view.focus.flow:
@@ -65,19 +70,17 @@ def set_focus(self, index):
         if self.master.commands.execute("view.properties.inbounds %d" % index):
             self.master.view.focus.index = index
 
-    def get_next(self, pos):
-        pos = pos + 1
-        if not self.master.commands.execute("view.properties.inbounds %d" % pos):
+    @lru_cache(maxsize=None)
+    def _get(self, pos: int) -> typing.Tuple[typing.Optional[FlowItem], typing.Optional[int]]:
+        if not self.master.view.inbounds(pos):
             return None, None
-        f = FlowItem(self.master, self.master.view[pos])
-        return f, pos
+        return FlowItem(self.master, self.master.view[pos]), pos
+
+    def get_next(self, pos):
+        return self._get(pos + 1)
 
     def get_prev(self, pos):
-        pos = pos - 1
-        if not self.master.commands.execute("view.properties.inbounds %d" % pos):
-            return None, None
-        f = FlowItem(self.master, self.master.view[pos])
-        return f, pos
+        return self._get(pos - 1)
 
 
 class FlowListBox(urwid.ListBox, layoutwidget.LayoutWidget):


Test patch: