Merge pull request #1815 from matrix-org/erikj/iter_cache_size

Optionally measure size of cache by sum of length of values

1 files changed, 16 insertions, 1 deletions

diff --git a/synapse/util/caches/treecache.py b/synapse/util/caches/treecache.py
index c31585aea3..fcc341a6b7 100644
--- a/synapse/util/caches/treecache.py
+++ b/synapse/util/caches/treecache.py
@@ -65,12 +65,27 @@ class TreeCache(object):
         return popped
 
     def values(self):
-        return [e.value for e in self.root.values()]
+        return list(iterate_tree_cache_entry(self.root))
 
     def __len__(self):
         return self.size
 
 
+def iterate_tree_cache_entry(d):
+    """Helper function to iterate over the leaves of a tree, i.e. a dict of that
+    can contain dicts.
+    """
+    if isinstance(d, dict):
+        for value_d in d.itervalues():
+            for value in iterate_tree_cache_entry(value_d):
+                yield value
+    else:
+        if isinstance(d, _Entry):
+            yield d.value
+        else:
+            yield d
+
+
 class _Entry(object):
     __slots__ = ["value"]