From c4664de41f354351e6c63efb0fa44a1e3905692a Mon Sep 17 00:00:00 2001
From: Georg Krause <mail@georg-krause.net>
Date: Thu, 8 Apr 2021 15:03:42 +0200
Subject: [PATCH] Add Middleware to trace memory usage

---
 .env.dev                               |  1 +
 api/config/settings/local.py           |  1 +
 api/funkwhale_api/common/middleware.py | 18 ++++++++++++++++++
 3 files changed, 20 insertions(+)

diff --git a/.env.dev b/.env.dev
index 40774610b..75ae332dc 100644
--- a/.env.dev
+++ b/.env.dev
@@ -13,6 +13,7 @@ BROWSABLE_API_ENABLED=True
 FORWARDED_PROTO=http
 LDAP_ENABLED=False
 FUNKWHALE_SPA_HTML_ROOT=http://nginx/front/
+PYTHONTRACEMALLOC=1
 
 # Uncomment this if you're using traefik/https
 # FORCE_HTTPS_URLS=True
diff --git a/api/config/settings/local.py b/api/config/settings/local.py
index b51ec273e..2fc121612 100644
--- a/api/config/settings/local.py
+++ b/api/config/settings/local.py
@@ -104,4 +104,5 @@ if env.bool("WEAK_PASSWORDS", default=False):
 MIDDLEWARE = (
     "funkwhale_api.common.middleware.DevHttpsMiddleware",
     "funkwhale_api.common.middleware.ProfilerMiddleware",
+    "funkwhale_api.common.middleware.PymallocMiddleware",
 ) + MIDDLEWARE
diff --git a/api/funkwhale_api/common/middleware.py b/api/funkwhale_api/common/middleware.py
index 475215910..f7d2acda4 100644
--- a/api/funkwhale_api/common/middleware.py
+++ b/api/funkwhale_api/common/middleware.py
@@ -14,6 +14,7 @@ from django.middleware import csrf
 from django.contrib import auth
 from django import urls
 from rest_framework import views
+import tracemalloc
 
 from funkwhale_api.federation import utils as federation_utils
 
@@ -405,3 +406,20 @@ class ProfilerMiddleware:
         response = http.HttpResponse("<pre>%s</pre>" % stream.getvalue())
 
         return response
+
+
+class PymallocMiddleware:
+    def __init__(self, get_response):
+        self.get_response = get_response
+
+    def __call__(self, request):
+
+        if tracemalloc.is_tracing():
+            snapshot = tracemalloc.take_snapshot()
+            stats = snapshot.statistics("lineno")
+
+            print("Memory trace")
+            for stat in stats[:25]:
+                print(stat)
+
+        return self.get_response(request)