optional memray

Author: JoshuaFox

config.yaml.original

@@ -66,3 +66,8 @@ test_or_dev_project_markers:
   - test
   - dev
   - qa
+
+# If True, the application will generate memray files downloadable with the memray_download tool.
+# Default false
+use_memray:
+    False

main.py

@@ -20,7 +20,7 @@
 
     app.wsgi_app = google.appengine.api.wrap_wsgi_app(app.wsgi_app)
 
-from util.utils import init_logging, tmp_dir
+from util.utils import init_logging, memray_tempdir
 
 # Must init logging before any library code writes logs (which would then just override our config)
 init_logging()
@@ -54,6 +54,7 @@
     is_in_test_or_dev_project,
     is_test_or_dev_configuration,
     iris_homepage_text,
+    use_memray,
 )
 from util.utils import log_time, timing
 
@@ -70,17 +71,15 @@
 logging.info(
     "env GAE_USE_SOCKETS_HTTPLIB is %s", os.environ.get("GAE_USE_SOCKETS_HTTPLIB")
 )
-logging.info(
-    "env PYTHONMALLOC is %s", os.environ.get("PYTHONMALLOC")
-)
+logging.info("env PYTHONMALLOC is %s", os.environ.get("PYTHONMALLOC"))
 
 PluginHolder.init()
 
 
 @app.route("/list_memray")
 def list_memray():
     s = ""
-    for f in os.listdir(tmp_dir()):
+    for f in os.listdir(memray_tempdir()):
         s += f"<a href='/memrayfiles/{f}'>{f}</a><br>\n"
     if not s:
         s = "NONE FOUND"
@@ -90,15 +89,15 @@ def list_memray():
 @app.route("/memrayfiles/<filename>")
 def memrayfiles(filename):
 
-    with open(f"{tmp_dir()}/{filename}", "rb") as f:
+    with open(f"{memray_tempdir()}/{filename}", "rb") as f:
         content = f.read()
 
     return Response(content, mimetype="application/octet-stream", status=200)
 
 
 @app.route("/")
 def index():
-    with memray.Tracker(memray_filename("index")):
+    def __index():
         increment_invocation_count("index")
         with gae_memory_logging("index"):
             msg = iris_homepage_text()
@@ -110,24 +109,38 @@ def index():
             )
             return Response(msg, mimetype="text/plain", status=200)
 
+    if use_memray():
+        with memray.Tracker(memray_filename("index")):
+            return __index()
+    else:
+        return __index()
+
 
 @app.route("/_ah/warmup")
 def warmup():
-    with memray.Tracker(memray_filename("warmup")):
+    def __warmup():
+
         increment_invocation_count("warmup")
         with gae_memory_logging("warmup"):
             logging.info("warmup() called")
 
         return "", 200, {}
 
+    if use_memray():
+        with memray.Tracker(memray_filename("warmup")):
+            return __warmup()
+    else:
+        return __warmup()
+
 
 @app.route("/schedule", methods=["GET"])
 @log_time
 def schedule():
     """
     Send out a message per-plugin per-project to label all objects of that type and project.
     """
-    with memray.Tracker(memray_filename("schedule")):
+
+    def __schedule():
 
         increment_invocation_count("schedule")
         with gae_memory_logging("schedule"):
@@ -147,6 +160,12 @@ def schedule():
                 logging.exception("In schedule()")
                 return "Error", 500
 
+    if use_memray():
+        with memray.Tracker(memray_filename("schedule")):
+            return __schedule()
+    else:
+        return __schedule()
+
 
 @lru_cache(maxsize=1)
 def __get_enabled_projects():
@@ -218,7 +237,7 @@ def __send_pubsub_per_projectplugin(configured_projects):
 
 @app.route("/label_one", methods=["POST"])
 def label_one():
-    with memray.Tracker(memray_filename("label_one")):
+    def __label_one():
 
         increment_invocation_count("label_one")
         with gae_memory_logging("label_one"):
@@ -275,6 +294,12 @@ def label_one():
                 logging.exception("Error on label_one %s %s", plugins_found, project_id)
                 return "Error", 500
 
+    if use_memray():
+        with memray.Tracker(memray_filename("label_one")):
+            return __label_one()
+    else:
+        return __label_one()
+
 
 def __label_one_0(data, plugin_cls: Type[Plugin]):
     plugin = PluginHolder.get_plugin_instance(plugin_cls)
@@ -330,7 +355,8 @@ def __extract_pubsub_content() -> Dict:
 
 @app.route("/do_label", methods=["POST"])
 def do_label():
-    with memray.Tracker(memray_filename("do_label")):
+    def __do_label():
+
         increment_invocation_count("do_label")
         with gae_memory_logging("do_label"):
 
@@ -373,6 +399,12 @@ def do_label():
                 )
                 return "Error", 500
 
+    if use_memray():
+        with memray.Tracker(memray_filename("do_label")):
+            return __do_label()
+    else:
+        return __do_label()
+
 
 def __check_pubsub_verification_token():
     """Token verifying that only PubSub accesses PubSub push endpoints"""

util/config_utils.py

@@ -3,6 +3,7 @@
 import re
 import sys
 import typing
+from collections import defaultdict
 
 import yaml
 
@@ -107,6 +108,10 @@ def is_test_or_dev_configuration():
     return get_config()["config_file"] != "config.yaml"
 
 
+def use_memray():
+    return get_config().get("use_memray")
+
+
 def is_in_test_or_dev_project(project_id):
     markers = get_config().get("test_or_dev_project_markers", [])
     for marker in markers:

util/gcp_utils.py

@@ -12,7 +12,7 @@
 
 from util import localdev_config, utils
 from util.detect_gae import detect_gae
-from util.utils import timed_lru_cache, log_time, dict_to_camelcase, tmp_dir
+from util.utils import timed_lru_cache, log_time, dict_to_camelcase, memray_tempdir
 
 __invocation_count = Counter()
 
@@ -240,7 +240,7 @@ def memray_filename(sample_loc):
         mem = ""
     global global_counter
     global_counter += 1
-    ret = f"{tmp_dir()}/{__inst_id}_{__isonow_for_filename()}.{global_counter}_{sample_loc}{mem}.bin"
+    ret = f"{memray_tempdir()}/{__inst_id}_{__isonow_for_filename()}.{global_counter}_{sample_loc}{mem}.bin"
     return ret
 
 

util/memray_download.py

@@ -13,7 +13,8 @@
 
 
 def download(base, path):
-    if base.endswith("/") and path.startswith("/"): path=path[1:]
+    if base.endswith("/") and path.startswith("/"):
+        path = path[1:]
     url = base + path
     try:
 
@@ -46,10 +47,10 @@ def download_list_and_bins(base):
     html = download_txt(base, "list_memray")
     parser = ExtractHrefs()
     parser.feed(html)
-    existing_bin=[Path(p).stem for p in Path(dir_with_bin()).iterdir()]
+    existing_bin = [Path(p).stem for p in Path(dir_with_bin()).iterdir()]
     for path in parser.paths:
         bin_pth = Path(path)
-        exist= bin_pth.stem in existing_bin
+        exist = bin_pth.stem in existing_bin
 
         if not exist:
             bin_content = download(base, path)
@@ -70,7 +71,9 @@ def generate_flamegraphs():
     mkdirs(dir_with_html_leaks())
     mkdirs(dir_with_html_basic())
 
-    binfiles = list(filter(lambda filename: filename.endswith(".bin"), os.listdir(dir_with_bin())))
+    binfiles = list(
+        filter(lambda filename: filename.endswith(".bin"), os.listdir(dir_with_bin()))
+    )
     print(len(binfiles), "bin files in", dir_with_bin())
 
     for f in binfiles:
@@ -79,7 +82,11 @@ def generate_flamegraphs():
             binfn = f"{dir_with_bin()}/{f}"
             sz = os.path.getsize(binfn)
             if sz > 1000:
-                convert_bin_to_one_html(binfn, stem_for_html, False, )
+                convert_bin_to_one_html(
+                    binfn,
+                    stem_for_html,
+                    False,
+                )
                 convert_bin_to_one_html(binfn, stem_for_html, True)
 
     count_html = sum(1 for _ in Path(dir_with_html_basic()).iterdir())
@@ -124,8 +131,7 @@ def main():
     base = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:8000"
     if not base.endswith("/"):
         base += "/"
-    MINUTES = 60
-    while time.time() - start < 10 * MINUTES:
+    while time.time() - start < 10 * 60:
         try:
             download_and_convert(base)
         except HTTPError as e:

util/utils.py

@@ -254,12 +254,11 @@ def run_command(command_s):
     return output.strip("\n")
 
 
-def tmp_dir() -> str:
-    subdir = "memray"
-    ret = f"/tmp/{subdir}"
+def memray_tempdir() -> str:
+    ret = f"/tmp/memray"
     mkdirs(ret)
     return ret
 
 
-def mkdirs(ret):
-    pathlib.Path(ret).mkdir(parents=True, exist_ok=True)
+def mkdirs(dir_):
+    pathlib.Path(dir_).mkdir(parents=True, exist_ok=True)