console: add caching sort keys

This is a tad complicated. The underlying sorted list implementation expects
the sort key to be stable for the lifetime of the object. However, if we sort
by size, for instance, the sort order changes as the flow progresses through
its lifecycle. We address this through two means:

- Let order keys cache the sort value by flow ID.
- Add a facility to refresh items in the list by removing and re-adding them
when they are updated.

3 files changed, 134 insertions, 70 deletions

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index 9d38d94c..c09ff454 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -19,39 +19,83 @@ import mitmproxy.flow
 from mitmproxy import flowfilter
 from mitmproxy import exceptions
 
+# The underlying sorted list implementation expects the sort key to be stable
+# for the lifetime of the object. However, if we sort by size, for instance,
+# the sort order changes as the flow progresses through its lifecycle. We
+# address this through two means:
+#
+# - Let order keys cache the sort value by flow ID.
+#
+# - Add a facility to refresh items in the list by removing and re-adding them
+# when they are updated.
+
+
+class _OrderKey:
+    def __init__(self, view):
+        self.view = view
 
-def key_request_start(f: mitmproxy.flow.Flow) -> datetime.datetime:
-    return f.request.timestamp_start or 0
+    def generate(self, f: mitmproxy.flow.Flow) -> typing.Any:
+        pass
 
+    def refresh(self, f):
+        k = self._key()
+        old = self.view.settings[f][k]
+        new = self.generate(f)
+        if old != new:
+            self.view._view.remove(f)
+            self.view.settings[f][k] = new
+            self.view._view.add(f)
+            self.view.sig_refresh.send(self.view)
 
-def key_request_method(f: mitmproxy.flow.Flow) -> str:
-    return f.request.method
+    def _key(self):
+        return "_order_%s" % id(self)
 
+    def __call__(self, f):
+        k = self._key()
+        s = self.view.settings[f]
+        if k in s:
+            return s[k]
+        val = self.generate(f)
+        s[k] = val
+        return val
 
-def key_request_url(f: mitmproxy.flow.Flow) -> str:
-    return f.request.url
 
+class OrderRequestStart(_OrderKey):
+    def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime:
+        return f.request.timestamp_start or 0
 
-def key_size(f: mitmproxy.flow.Flow) -> int:
-    s = 0
-    if f.request.raw_content:
-        s += len(f.request.raw_content)
-    if f.response and f.response.raw_content:
-        s += len(f.response.raw_content)
-    return s
 
+class OrderRequestMethod(_OrderKey):
+    def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime:
+        return f.request.method
 
-orders = [
-    ("t", "time", key_request_start),
-    ("m", "method", key_request_method),
-    ("u", "url", key_request_url),
-    ("z", "size", key_size),
-]
+
+class OrderRequestURL(_OrderKey):
+    def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime:
+        return f.request.url
+
+
+class OrderKeySize(_OrderKey):
+    def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime:
+        s = 0
+        if f.request.raw_content:
+            s += len(f.request.raw_content)
+        if f.response and f.response.raw_content:
+            s += len(f.response.raw_content)
+        return s
 
 
 matchall = flowfilter.parse(".")
 
 
+orders = [
+    ("t", "time"),
+    ("m", "method"),
+    ("u", "url"),
+    ("z", "size"),
+]
+
+
 class View(collections.Sequence):
     def __init__(self):
         super().__init__()
@@ -59,7 +103,15 @@ class View(collections.Sequence):
         self.filter = matchall
         # Should we show only marked flows?
         self.show_marked = False
-        self.order_key = key_request_start
+
+        self.default_order = OrderRequestStart(self)
+        self.orders = dict(
+            time = self.default_order,
+            method = OrderRequestMethod(self),
+            url = OrderRequestURL(self),
+            size = OrderKeySize(self),
+        )
+        self.order_key = self.default_order
         self.order_reversed = False
         self.focus_follow = False
 
@@ -120,13 +172,20 @@ class View(collections.Sequence):
     def __contains__(self, f: mitmproxy.flow.Flow) -> bool:
         return self._view.__contains__(f)
 
+    def _order_key_name(self):
+        return "_order_%s" % id(self.order_key)
+
+    def _base_add(self, f):
+        self.settings[f][self._order_key_name()] = self.order_key(f)
+        self._view.add(f)
+
     def _refilter(self):
         self._view.clear()
         for i in self._store.values():
             if self.show_marked and not i.marked:
                 continue
             if self.filter(i):
-                self._view.add(i)
+                self._base_add(i)
         self.sig_refresh.send(self)
 
     # API
@@ -160,9 +219,10 @@ class View(collections.Sequence):
         """
         self._store.clear()
         self._view.clear()
+        self.settings.clear()
         self.sig_refresh.send(self)
 
-    def add(self, f: mitmproxy.flow.Flow):
+    def add(self, f: mitmproxy.flow.Flow) -> bool:
         """
             Adds a flow to the state. If the flow already exists, it is
             ignored.
@@ -170,7 +230,7 @@ class View(collections.Sequence):
         if f.id not in self._store:
             self._store[f.id] = f
             if self.filter(f):
-                self._view.add(f)
+                self._base_add(f)
                 if self.focus_follow:
                     self.focus.flow = f
                 self.sig_add.send(self, flow=f)
@@ -180,10 +240,10 @@ class View(collections.Sequence):
             Removes the flow from the underlying store and the view.
         """
         if f.id in self._store:
-            del self._store[f.id]
             if f in self._view:
                 self._view.remove(f)
                 self.sig_remove.send(self, flow=f)
+            del self._store[f.id]
 
     def update(self, f: mitmproxy.flow.Flow):
         """
@@ -192,11 +252,16 @@ class View(collections.Sequence):
         if f.id in self._store:
             if self.filter(f):
                 if f not in self._view:
-                    self._view.add(f)
+                    self._base_add(f)
                     if self.focus_follow:
                         self.focus.flow = f
                     self.sig_add.send(self, flow=f)
                 else:
+                    # This is a tad complicated. The sortedcontainers
+                    # implementation assumes that the order key is stable. If
+                    # it changes mid-way Very Bad Things happen. We detect when
+                    # this happens, and re-fresh the item.
+                    self.order_key.refresh(f)
                     self.sig_update.send(self, flow=f)
             else:
                 try:
@@ -219,16 +284,13 @@ class View(collections.Sequence):
             self.set_filter(filt)
         if "order" in updated:
             if opts.order is None:
-                self.set_order(key_request_start)
+                self.set_order(self.default_order)
             else:
-                for _, name, func in orders:
-                    if name == opts.order:
-                        self.set_order(func)
-                        break
-                else:
+                if opts.order not in self.orders:
                     raise exceptions.OptionsError(
                         "Unknown flow order: %s" % opts.order
                     )
+                self.set_order(self.orders[opts.order])
         if "order_reversed" in updated:
             self.set_reversed(opts.order_reversed)
         if "focus_follow" in updated:
@@ -237,16 +299,16 @@ class View(collections.Sequence):
     def request(self, f):
         self.add(f)
 
-    def intercept(self, f):
+    def error(self, f):
         self.update(f)
 
-    def resume(self, f):
+    def response(self, f):
         self.update(f)
 
-    def error(self, f):
+    def intercept(self, f):
         self.update(f)
 
-    def response(self, f):
+    def resume(self, f):
         self.update(f)
 
 
@@ -316,6 +378,9 @@ class Settings(collections.Mapping):
         view.sig_remove.connect(self._sig_remove)
         view.sig_refresh.connect(self._sig_refresh)
 
+    def clear(self):
+        self.values.clear()
+
     def __iter__(self) -> typing.Iterable:
         return iter(self.values)
 
@@ -332,6 +397,6 @@ class Settings(collections.Mapping):
             del self.values[flow.id]
 
     def _sig_refresh(self, view):
-        for fid in self.values.keys():
+        for fid in list(self.values.keys()):
             if fid not in view._store:
                 del self.values[fid]
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 31ce17a3..ffbfb0cb 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -92,10 +92,14 @@ class Master:
         try:
             mtype, obj = self.event_queue.get(timeout=timeout)
             if mtype not in events.Events:
-                raise exceptions.ControlException("Unknown event %s" % repr(mtype))
+                raise exceptions.ControlException(
+                    "Unknown event %s" % repr(mtype)
+                )
             handle_func = getattr(self, mtype)
             if not callable(handle_func):
-                raise exceptions.ControlException("Handler %s not callable" % mtype)
+                raise exceptions.ControlException(
+                    "Handler %s not callable" % mtype
+                )
             if not handle_func.__dict__.get("__handler"):
                 raise exceptions.ControlException(
                     "Handler function %s is not decorated with controller.handler" % (
diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py
index 77a8da30..63df8307 100644
--- a/test/mitmproxy/addons/test_view.py
+++ b/test/mitmproxy/addons/test_view.py
@@ -7,9 +7,18 @@ from mitmproxy.test import taddons
 from .. import tutils
 
 
-def test_keys():
-    t = tflow.tflow(resp=True)
-    view.key_size(t)
+class Options(options.Options):
+    def __init__(
+        self, *,
+        filter=None,
+        order=None,
+        order_reversed=False,
+        **kwargs
+    ):
+        self.filter = filter
+        self.order = order
+        self.order_reversed = order_reversed
+        super().__init__(**kwargs)
 
 
 def test_simple():
@@ -74,22 +83,23 @@ def test_filter():
 
 def test_order():
     v = view.View()
-    v.request(tft(method="get", start=1))
-    v.request(tft(method="put", start=2))
-    v.request(tft(method="get", start=3))
-    v.request(tft(method="put", start=4))
-    assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4]
-
-    v.set_order(view.key_request_method)
-    assert [i.request.method for i in v] == ["GET", "GET", "PUT", "PUT"]
-    v.set_reversed(True)
-    assert [i.request.method for i in v] == ["PUT", "PUT", "GET", "GET"]
+    with taddons.context(options=Options()) as tctx:
+        v.request(tft(method="get", start=1))
+        v.request(tft(method="put", start=2))
+        v.request(tft(method="get", start=3))
+        v.request(tft(method="put", start=4))
+        assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4]
+
+        tctx.configure(v, order="method")
+        assert [i.request.method for i in v] == ["GET", "GET", "PUT", "PUT"]
+        v.set_reversed(True)
+        assert [i.request.method for i in v] == ["PUT", "PUT", "GET", "GET"]
 
-    v.set_order(view.key_request_start)
-    assert [i.request.timestamp_start for i in v] == [4, 3, 2, 1]
+        tctx.configure(v, order="time")
+        assert [i.request.timestamp_start for i in v] == [4, 3, 2, 1]
 
-    v.set_reversed(False)
-    assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4]
+        v.set_reversed(False)
+        assert [i.request.timestamp_start for i in v] == [1, 2, 3, 4]
 
 
 def test_reversed():
@@ -261,7 +271,6 @@ def test_settings():
 
     tutils.raises(KeyError, v.settings.__getitem__, f)
     v.add(f)
-    assert v.settings[f] == {}
     v.settings[f]["foo"] = "bar"
     assert v.settings[f]["foo"] == "bar"
     assert len(list(v.settings)) == 1
@@ -270,20 +279,6 @@ def test_settings():
     assert not v.settings.keys()
 
 
-class Options(options.Options):
-    def __init__(
-        self, *,
-        filter=None,
-        order=None,
-        order_reversed=False,
-        **kwargs
-    ):
-        self.filter = filter
-        self.order = order
-        self.order_reversed = order_reversed
-        super().__init__(**kwargs)
-
-
 def test_configure():
     v = view.View()
     with taddons.context(options=Options()) as tctx: