session: implemented changes requested after PR review.

mitmproxy/addons/session.py

@@ -50,12 +50,13 @@ class SessionDB:
         or create a new one with optional path.
         :param db_path:
         """
-        self.live_components = {}
+        self.live_components: typing.Dict[str, tuple] = {}
-        self.tempdir = None
+        self.tempdir: tempfile.TemporaryDirectory = None
-        self.con = None
+        self.con: sqlite3.Connection = None
         # This is used for fast look-ups over bodies already dumped to database.
         # This permits to enforce one-to-one relationship between flow and body table.
-        self.body_ledger = set()
+        self.body_ledger: typing.Set[str] = set()
+        self.id_ledger: typing.Set[str] = set()
         if db_path is not None and os.path.isfile(db_path):
             self._load_session(db_path)
         else:
@@ -74,14 +75,10 @@ class SessionDB:
             shutil.rmtree(self.tempdir)
 
     def __contains__(self, fid):
-        return any([fid == i for i in self._get_ids()])
+        return fid in self.id_ledger
 
     def __len__(self):
-        ln = self.con.execute("SELECT COUNT(*) FROM flow;").fetchall()[0]
+        return len(self.id_ledger)
-        return ln[0] if ln else 0
-
-    def _get_ids(self):
-        return [t[0] for t in self.con.execute("SELECT id FROM flow;").fetchall()]
 
     def _load_session(self, path):
         if not self.is_session_db(path):
@@ -150,6 +147,7 @@ class SessionDB:
         body_buf = []
         flow_buf = []
         for flow in flows:
+            self.id_ledger.add(flow.id)
             self._disassemble(flow)
             f = copy.copy(flow)
             f.request = copy.deepcopy(flow.request)
@@ -209,17 +207,21 @@ orders = [
 
 
 class Session:
+
+    _FP_RATE = 150
+    _FP_DECREMENT = 0.9
+    _FP_DEFAULT = 3.0
+
     def __init__(self):
-        self.db_store = None
+        self.db_store: SessionDB = None
-        self._hot_store = collections.OrderedDict()
+        self._hot_store: collections.OrderedDict = collections.OrderedDict()
-        self._live_components = {}
+        self._order_store: typing.Dict[str, typing.Dict[str, typing.Union[int, float, str]]] = {}
-        self._view = []
+        self._view: typing.List[typing.Tuple[typing.Union[int, float, str], str]] = []
-        self.order = orders[0]
+        self.order: str = orders[0]
         self.filter = matchall
-        self._flush_period = 3.0
+        self._flush_period: float = self._FP_DEFAULT
-        self._tweak_period = 0.5
+        self._flush_rate: int = self._FP_RATE
-        self._flush_rate = 150
+        self.started: bool = False
-        self.started = False
 
     def load(self, loader):
         loader.add_option(
@@ -242,7 +244,6 @@ class Session:
             self.db_store = SessionDB(ctx.options.session_path)
             loop = asyncio.get_event_loop()
             loop.create_task(self._writer())
-            loop.create_task(self._tweaker())
 
     def configure(self, updated):
         if "view_order" in updated:
@@ -253,28 +254,23 @@ class Session:
     async def _writer(self):
         while True:
             await asyncio.sleep(self._flush_period)
-            tof = []
+            batches = -(-len(self._hot_store) // self._flush_rate)
-            to_dump = min(self._flush_rate, len(self._hot_store))
+            self._flush_period = self._flush_period * self._FP_DECREMENT if batches > 1 else self._FP_DEFAULT
-            for _ in range(to_dump):
+            while batches:
-                tof.append(self._hot_store.popitem(last=False)[1])
+                tof = []
-            self.db_store.store_flows(tof)
+                to_dump = min(len(self._hot_store), self._flush_rate)
+                for _ in range(to_dump):
+                    tof.append(self._hot_store.popitem(last=False)[1])
+                self.db_store.store_flows(tof)
+                batches -= 1
+                await asyncio.sleep(0.01)
 
-    async def _tweaker(self):
+    def load_view(self) -> typing.Sequence[http.HTTPFlow]:
-        while True:
-            await asyncio.sleep(self._tweak_period)
-            if len(self._hot_store) >= 3 * self._flush_rate:
-                self._flush_period *= 0.9
-                self._flush_rate *= 1.1
-            elif len(self._hot_store) < self._flush_rate:
-                self._flush_period *= 1.1
-                self._flush_rate *= 0.9
-
-    def load_view(self):
         ids = [fid for _, fid in self._view]
         flows = self.load_storage(ids)
-        return sorted(flows, key=lambda f: self._generate_order(f))
+        return sorted(flows, key=lambda f: self._generate_order(self.order, f))
 
-    def load_storage(self, ids=None):
+    def load_storage(self, ids=None) -> typing.Sequence[http.HTTPFlow]:
         flows = []
         ids_from_store = []
         if ids is not None:
@@ -284,8 +280,6 @@ class Session:
                     flows.append(self._hot_store[fid])
                 elif fid in self.db_store:
                     ids_from_store.append(fid)
-                else:
-                    flows.append(None)
             flows += self.db_store.retrieve_flows(ids_from_store)
         else:
             for flow in self._hot_store.values():
@@ -300,15 +294,15 @@ class Session:
         self._hot_store.clear()
         self._view = []
 
-    def store_count(self):
+    def store_count(self) -> int:
         ln = 0
         for fid in self._hot_store.keys():
             if fid not in self.db_store:
                 ln += 1
         return ln + len(self.db_store)
 
-    def _generate_order(self, f: http.HTTPFlow) -> typing.Optional[typing.Union[str, int, float]]:
+    @staticmethod
-        o = self.order
+    def _generate_order(o: str, f: http.HTTPFlow) -> typing.Optional[typing.Union[str, int, float]]:
         if o == "time":
             return f.request.timestamp_start or 0
         if o == "method":
@@ -324,6 +318,11 @@ class Session:
             return s
         return None
 
+    def _store_order(self, f: http.HTTPFlow):
+        self._order_store[f.id] = {}
+        for order in orders:
+            self._order_store[f.id][order] = self._generate_order(order, f)
+
     def set_order(self, order: str) -> None:
         if order not in orders:
             raise CommandError(
@@ -332,7 +331,7 @@ class Session:
         if order != self.order:
             self.order = order
             newview = [
-                (self._generate_order(f), f.id) for f in self.load_view()
+                (self._order_store[t[1]][order], t[1]) for t in self._view
             ]
             self._view = sorted(newview)
 
@@ -341,7 +340,7 @@ class Session:
         flows = self.load_storage()
         for f in flows:
             if self.filter(f):
-                self._base_add(f)
+                self.update_view(f)
 
     def set_filter(self, input_filter: typing.Optional[str]) -> None:
         filt = matchall if not input_filter else flowfilter.parse(input_filter)
@@ -352,22 +351,20 @@ class Session:
         self.filter = filt
         self._refilter()
 
-    def _base_add(self, f):
+    def update_view(self, f):
-        if not any([f.id == t[1] for t in self._view]):
+        if any([f.id == t[1] for t in self._view]):
-            o = self._generate_order(f)
-            self._view.insert(bisect.bisect_left(KeyifyList(self._view, lambda x: x[0]), o), (o, f.id))
-        else:
-            o = self._generate_order(f)
             self._view = [(order, fid) for order, fid in self._view if fid != f.id]
-            self._view.insert(bisect.bisect_left(KeyifyList(self._view, lambda x: x[0]), o), (o, f.id))
+        o = self._order_store[f.id][self.order]
+        self._view.insert(bisect.bisect_left(KeyifyList(self._view, lambda x: x[0]), o), (o, f.id))
 
     def update(self, flows: typing.Sequence[http.HTTPFlow]) -> None:
         for f in flows:
+            self._store_order(f)
             if f.id in self._hot_store:
                 self._hot_store.pop(f.id)
             self._hot_store[f.id] = f
             if self.filter(f):
-                self._base_add(f)
+                self.update_view(f)
 
     def request(self, f):
         self.update([f])

test/mitmproxy/addons/test_session.py

@@ -28,8 +28,10 @@ class TestSession:
             tctx.master.addons.add(s)
             tctx.options.session_path = None
             tctx.options.view_filter = None
+        # To make tests quicker
         if fp:
             s._flush_period = fp
+            s._FP_DEFAULT = fp
         s.running()
         return s
 
@@ -85,21 +87,11 @@ class TestSession:
     def test_session_order_generators(self):
         s = session.Session()
         tf = tflow.tflow(resp=True)
-
+        assert s._generate_order('time', tf) == 946681200
-        s.order = "time"
+        assert s._generate_order('method', tf) == tf.request.method
-        assert s._generate_order(tf) == 946681200
+        assert s._generate_order('url', tf) == tf.request.url
-
+        assert s._generate_order('size', tf) == len(tf.request.raw_content) + len(tf.response.raw_content)
-        s.order = "method"
+        assert not s._generate_order('invalid', tf)
-        assert s._generate_order(tf) == tf.request.method
-
-        s.order = "url"
-        assert s._generate_order(tf) == tf.request.url
-
-        s.order = "size"
-        assert s._generate_order(tf) == len(tf.request.raw_content) + len(tf.response.raw_content)
-
-        s.order = "invalid"
-        assert not s._generate_order(tf)
 
     def test_storage_simple(self):
         s = session.Session()
@@ -110,8 +102,12 @@ class TestSession:
         assert s.store_count() == 0
         s.request(f)
         assert s._view == [(1, f.id)]
+        assert s._order_store[f.id]['time'] == 1
+        assert s._order_store[f.id]['method'] == f.request.method
+        assert s._order_store[f.id]['url'] == f.request.url
+        assert s._order_store[f.id]['size'] == len(f.request.raw_content)
         assert s.load_view() == [f]
-        assert s.load_storage(['nonexistent']) == [None]
+        assert s.load_storage(['nonexistent']) == []
 
         s.error(f)
         s.response(f)
@@ -121,6 +117,10 @@ class TestSession:
 
         # Verify that flow has been updated, not duplicated
         assert s._view == [(1, f.id)]
+        assert s._order_store[f.id]['time'] == 1
+        assert s._order_store[f.id]['method'] == f.request.method
+        assert s._order_store[f.id]['url'] == f.request.url
+        assert s._order_store[f.id]['size'] == len(f.request.raw_content)
         assert s.store_count() == 1
 
         f2 = self.tft(start=3)
@@ -174,16 +174,16 @@ class TestSession:
 
         f.server_conn.via = tflow.tserver_conn()
         s.request(f)
-        await asyncio.sleep(1)
+        await asyncio.sleep(0.6)
         assert len(s._hot_store) == 0
         assert all([lflow.__dict__ == flow.__dict__ for lflow, flow in list(zip(s.load_storage(), [f]))])
 
         flows = [self.tft() for _ in range(500)]
         s.update(flows)
-        fp = s._flush_period
-        fr = s._flush_rate
         await asyncio.sleep(0.6)
-        assert s._flush_period < fp and s._flush_rate > fr
+        assert s._flush_period == s._FP_DEFAULT * s._FP_DECREMENT
+        await asyncio.sleep(3)
+        assert s._flush_period == s._FP_DEFAULT
 
     @pytest.mark.asyncio
     async def test_storage_bodies(self):