From e9f5812efff29b08bd35724689d10a362b410168 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Fri, 23 Apr 2021 11:59:43 +0100
Subject: [PATCH] Track memory usage of caches

---
 mypy.ini                        |  3 ++
 synapse/python_dependencies.py  |  1 +
 synapse/util/caches/__init__.py |  8 +++++
 synapse/util/caches/lrucache.py | 52 ++++++++++++++++++++++++++++++++-
 4 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/mypy.ini b/mypy.ini
index 32e6197409..c46d5bb722 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -172,3 +172,6 @@ ignore_missing_imports = True
 
 [mypy-txacme.*]
 ignore_missing_imports = True
+
+[mypy-pympler.*]
+ignore_missing_imports = True
diff --git a/synapse/python_dependencies.py b/synapse/python_dependencies.py
index 2a1c925ee8..77176a6cd9 100644
--- a/synapse/python_dependencies.py
+++ b/synapse/python_dependencies.py
@@ -121,6 +121,7 @@ CONDITIONAL_REQUIREMENTS = {
     # hiredis is not a *strict* dependency, but it makes things much faster.
     # (if it is not installed, we fall back to slow code.)
     "redis": ["txredisapi>=1.4.7", "hiredis"],
+    "cache_memroy": ["pympler"],
 }
 
 ALL_OPTIONAL_REQUIREMENTS = set()  # type: Set[str]
diff --git a/synapse/util/caches/__init__.py b/synapse/util/caches/__init__.py
index 46af7fa473..8e253cb518 100644
--- a/synapse/util/caches/__init__.py
+++ b/synapse/util/caches/__init__.py
@@ -32,6 +32,11 @@ cache_hits = Gauge("synapse_util_caches_cache:hits", "", ["name"])
 cache_evicted = Gauge("synapse_util_caches_cache:evicted_size", "", ["name"])
 cache_total = Gauge("synapse_util_caches_cache:total", "", ["name"])
 cache_max_size = Gauge("synapse_util_caches_cache_max_size", "", ["name"])
+cache_memory_usage = Gauge(
+    "synapse_util_caches_cache_memory_usage",
+    "Estimated size in bytes of the caches",
+    ["name"],
+)
 
 response_cache_size = Gauge("synapse_util_caches_response_cache:size", "", ["name"])
 response_cache_hits = Gauge("synapse_util_caches_response_cache:hits", "", ["name"])
@@ -52,6 +57,7 @@ class CacheMetric:
     hits = attr.ib(default=0)
     misses = attr.ib(default=0)
     evicted_size = attr.ib(default=0)
+    memory_usage = attr.ib(default=None)
 
     def inc_hits(self):
         self.hits += 1
@@ -81,6 +87,8 @@ class CacheMetric:
                 cache_total.labels(self._cache_name).set(self.hits + self.misses)
                 if getattr(self._cache, "max_size", None):
                     cache_max_size.labels(self._cache_name).set(self._cache.max_size)
+                if self.memory_usage is not None:
+                    cache_memory_usage.labels(self._cache_name).set(self.memory_usage)
             if self._collect_callback:
                 self._collect_callback()
         except Exception as e:
diff --git a/synapse/util/caches/lrucache.py b/synapse/util/caches/lrucache.py
index a21d34fcb4..97cc77156a 100644
--- a/synapse/util/caches/lrucache.py
+++ b/synapse/util/caches/lrucache.py
@@ -33,6 +33,33 @@ from synapse.config import cache as cache_config
 from synapse.util.caches import CacheMetric, register_cache
 from synapse.util.caches.treecache import TreeCache
 
+try:
+    from pympler.asizeof import Asizer
+
+    sizer = Asizer()
+    sizer.exclude_refs((), None, "")
+
+    def _get_size_of(val: Any, *, recurse=True) -> int:
+        """Get an estimate of the size in bytes of the object.
+
+        Args:
+            val: The object to size.
+            recurse: If true will include referenced values in the size,
+                otherwise only sizes the given object.
+        """
+        return sizer.asizeof(val, limit=100 if recurse else 0)
+
+
+except ImportError:
+
+    def _get_size_of(val: Any, recurse=True) -> int:
+        return 0
+
+
+# Whether to track estimated memory usage of the LruCaches.
+TRACK_MEMORY_USAGE = True
+
+
 # Function type: the type used for invalidation callbacks
 FT = TypeVar("FT", bound=Callable[..., Any])
 
@@ -54,7 +81,7 @@ def enumerate_leaves(node, depth):
 
 
 class _Node:
-    __slots__ = ["prev_node", "next_node", "key", "value", "callbacks"]
+    __slots__ = ["prev_node", "next_node", "key", "value", "callbacks", "memory"]
 
     def __init__(
         self, prev_node, next_node, key, value, callbacks: Optional[set] = None
@@ -65,6 +92,16 @@ class _Node:
         self.value = value
         self.callbacks = callbacks or set()
 
+        self.memory = 0
+        if TRACK_MEMORY_USAGE:
+            self.memory = (
+                _get_size_of(key)
+                + _get_size_of(value)
+                + _get_size_of(self.callbacks, recurse=False)
+                + _get_size_of(self, recurse=False)
+            )
+            self.memory += _get_size_of(self.memory, recurse=False)
+
 
 class LruCache(Generic[KT, VT]):
     """
@@ -136,6 +173,9 @@ class LruCache(Generic[KT, VT]):
                 self,
                 collect_callback=metrics_collection_callback,
             )  # type: Optional[CacheMetric]
+
+            if TRACK_MEMORY_USAGE and metrics:
+                metrics.memory_usage = 0
         else:
             metrics = None
 
@@ -188,6 +228,9 @@ class LruCache(Generic[KT, VT]):
             if size_callback:
                 cached_cache_len[0] += size_callback(node.value)
 
+            if TRACK_MEMORY_USAGE and metrics:
+                metrics.memory_usage += node.memory
+
         def move_node_to_front(node):
             prev_node = node.prev_node
             next_node = node.next_node
@@ -214,6 +257,10 @@ class LruCache(Generic[KT, VT]):
             for cb in node.callbacks:
                 cb()
             node.callbacks.clear()
+
+            if TRACK_MEMORY_USAGE and metrics:
+                metrics.memory_usage -= node.memory
+
             return deleted_len
 
         @overload
@@ -332,6 +379,9 @@ class LruCache(Generic[KT, VT]):
             if size_callback:
                 cached_cache_len[0] = 0
 
+            if TRACK_MEMORY_USAGE and metrics:
+                metrics.memory_usage = 0
+
         @synchronized
         def cache_contains(key: KT) -> bool:
             return key in cache