Add tracemalloc to stress_lowlevel.py

Author: benjeffery

python/stress_lowlevel.py

@@ -3,12 +3,13 @@
 any memory leaks or error handling issues.
 """
 import argparse
+import curses
 import logging
 import os
 import random
 import resource
-import sys
 import time
+import tracemalloc
 import unittest
 
 import tests.test_dict_encoding as test_dict_encoding
@@ -25,7 +26,7 @@
 import tests.test_vcf as test_vcf
 
 
-def main():
+def main(stdscr):
     modules = {
         "highlevel": test_highlevel,
         "lowlevel": test_lowlevel,
@@ -57,13 +58,14 @@ def main():
     # Need to do this to silence the errors from the file_format tests.
     logging.basicConfig(level=logging.ERROR)
 
-    print("iter\ttests\terr\tfail\tskip\tRSS\tmin\tmax\tmax@iter")
     max_rss = 0
     max_rss_iter = 0
     min_rss = 1e100
     iteration = 0
     last_print = time.time()
     devnull = open(os.devnull, "w")
+    tracemalloc.start()
+    memory_start = None
     while True:
         # We don't want any random variation in the amount of memory
         # used from test-to-test.
@@ -73,7 +75,10 @@ def main():
         for mod in test_modules[1:]:
             suite.addTests(testloader.loadTestsFromModule(mod))
         runner = unittest.TextTestRunner(verbosity=0, stream=devnull)
+        if memory_start is None:
+            memory_start = tracemalloc.take_snapshot()
         result = runner.run(suite)
+        memory_current = tracemalloc.take_snapshot()
         rusage = resource.getrusage(resource.RUSAGE_SELF)
         if max_rss < rusage.ru_maxrss:
             max_rss = rusage.ru_maxrss
@@ -83,24 +88,46 @@ def main():
 
         # We don't want to flood stdout, so we rate-limit to 1 per second.
         if time.time() - last_print > 1:
-            print(
-                iteration,
-                result.testsRun,
-                len(result.failures),
-                len(result.errors),
-                len(result.skipped),
-                rusage.ru_maxrss,
-                min_rss,
-                max_rss,
-                max_rss_iter,
-                sep="\t",
-                end="\r",
+            stdscr.clear()
+            stdscr.addstr(0, 0, "iter\ttests\terr\tfail\tskip\tRSS\tmin\tmax\tmax@iter")
+            stdscr.addstr(
+                1,
+                0,
+                "\t".join(
+                    map(
+                        str,
+                        [
+                            iteration,
+                            result.testsRun,
+                            len(result.failures),
+                            len(result.errors),
+                            len(result.skipped),
+                            rusage.ru_maxrss,
+                            min_rss,
+                            max_rss,
+                            max_rss_iter,
+                        ],
+                    )
+                ),
             )
+            stats = memory_current.compare_to(memory_start, "traceback")
+            rows, cols = stdscr.getmaxyx()
+            for i, stat in enumerate(stats[: rows - 3], 1):
+                stdscr.addstr(i + 2, 0, str(stat))
             last_print = time.time()
-            sys.stdout.flush()
+            stdscr.refresh()
 
         iteration += 1
 
 
 if __name__ == "__main__":
-    main()
+    stdscr = curses.initscr()
+    curses.noecho()
+    curses.cbreak()
+
+    try:
+        main(stdscr)
+    finally:
+        curses.echo()
+        curses.nocbreak()
+        curses.endwin()