From 75ccd44c42458eb9fafc87b66d7007cf5259f487 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Fri, 3 Sep 2021 21:22:33 +0200
Subject: [PATCH] add more benchmark/memoryleak tooling

---
 test/helper_tools/bench.sh             | 26 ++++++++
 test/helper_tools/hunt_memory_leaks.py | 87 ++++++++++++++++++++++++++
 2 files changed, 113 insertions(+)
 create mode 100644 test/helper_tools/bench.sh
 create mode 100644 test/helper_tools/hunt_memory_leaks.py

diff --git a/test/helper_tools/bench.sh b/test/helper_tools/bench.sh
new file mode 100644
index 000000000..cc48e4b7a
--- /dev/null
+++ b/test/helper_tools/bench.sh
@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+
+kill -s KILL $(pidof caddy)
+
+caddy file-server -listen 127.0.0.1:8000 &
+
+benchit() {
+  kill -s KILL $(pidof python3)
+  mitmdump -q &
+  sleep 10
+  echo "bench..."
+  hey -x http://127.0.0.1:8080 -n 2000 -disable-keepalive http://127.0.0.1:8000/bench.sh
+}
+
+for i in {1..2} ; do
+  git stash -q
+  benchit
+  echo "^ without changes"
+
+  git stash pop -q
+  benchit
+  echo "^ with changes"
+done
+
+kill -s KILL $(pidof caddy)
+kill -s KILL $(pidof python3)
diff --git a/test/helper_tools/hunt_memory_leaks.py b/test/helper_tools/hunt_memory_leaks.py
new file mode 100644
index 000000000..9be5db1c8
--- /dev/null
+++ b/test/helper_tools/hunt_memory_leaks.py
@@ -0,0 +1,87 @@
+import collections
+import gc
+import os
+import signal
+
+from mitmproxy import flow
+
+
+def load(loader):
+    signal.signal(signal.SIGUSR1, debug1)
+    signal.signal(signal.SIGUSR2, debug2)
+    print(f"Debug signal registered. Run the following commands for diagnostics:")
+    print()
+    print(f"  kill -s USR1 {os.getpid()}")
+    print(f"  kill -s USR2 {os.getpid()}")
+    print()
+
+
+def debug1(*_):
+    print()
+    print("Before GC")
+    print("=======")
+    print("gc.get_stats", gc.get_stats())
+    print("gc.get_count", gc.get_count())
+    print("gc.get_threshold", gc.get_threshold())
+
+    gc.collect()
+
+    print()
+    print("After GC")
+    print("=======")
+    print("gc.get_stats", gc.get_stats())
+    print("gc.get_count", gc.get_count())
+    print("gc.get_threshold", gc.get_threshold())
+
+    print()
+    print("Memory")
+    print("=======")
+    for t, count in collections.Counter([str(type(o)) for o in gc.get_objects()]).most_common(50):
+        print(count, t)
+
+
+def debug2(*_):
+    print()
+    print("Flow References")
+    print("=======")
+
+    #gc.collect()
+
+    objs = tuple(gc.get_objects())
+    ignore = {id(objs)}
+    flows = 0
+    for i in range(len(objs)):
+        try:
+            is_flow = isinstance(objs[i], flow.Flow)
+        except Exception:
+            continue
+        if is_flow:
+            flows += 1
+            #print_refs(objs[i], ignore, set())
+            #break
+    del objs
+
+    print(f"{flows} flows found.")
+
+
+def print_refs(x, ignore: set, seen: set, depth: int = 0, max_depth: int = 10):
+    if id(x) in ignore:
+        return
+
+    if id(x) in seen:
+        print("  " * depth + "↖ " + repr(str(x))[1:60] + f" (\x1b[31mseen\x1b[0m: {id(x):x})")
+        return
+    else:
+        if depth == 0:
+            print("- " + repr(str(x))[1:60] + f" ({id(x):x})")
+        else:
+            print("  " * depth + "↖ " + repr(str(x))[1:60] + f" ({id(x):x})")
+        seen.add(id(x))
+
+    if depth == max_depth:
+        return
+
+    referrers = tuple(gc.get_referrers(x))
+    ignore.add(id(referrers))
+    for ref in referrers:
+        print_refs(ref, ignore, seen, depth + 1, max_depth)