From 766526e1142e7ad0ffb43bd075b0ff2d6265e4cb Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Fri, 29 Jan 2016 10:11:21 +0000
Subject: Make TreeCache keep track of its own size.

---
 synapse/util/caches/treecache.py | 7 +++++++
 1 file changed, 7 insertions(+)

(limited to 'synapse/util/caches/treecache.py')

diff --git a/synapse/util/caches/treecache.py b/synapse/util/caches/treecache.py
index 3b58860910..a29ea8144e 100644
--- a/synapse/util/caches/treecache.py
+++ b/synapse/util/caches/treecache.py
@@ -8,6 +8,7 @@ class TreeCache(object):
     Keys must be tuples.
     """
     def __init__(self):
+        self.size = 0
         self.root = {}
 
     def __setitem__(self, key, value):
@@ -21,6 +22,7 @@ class TreeCache(object):
         for k in key[:-1]:
             node = node.setdefault(k, {})
         node[key[-1]] = value
+        self.size += 1
 
     def get(self, key, default=None):
         node = self.root
@@ -31,6 +33,7 @@ class TreeCache(object):
         return node.get(key[-1], default)
 
     def clear(self):
+        self.size = 0
         self.root = {}
 
     def pop(self, key, default=None):
@@ -57,4 +60,8 @@ class TreeCache(object):
                 break
             node_and_keys[i+1][0].pop(k)
 
+        self.size -= 1
         return popped
+
+    def __len__(self):
+        return self.size
-- 
cgit 1.4.1


From a30364c1f99bdd7d5cb0fe82ebdfe52d996defef Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Fri, 29 Jan 2016 10:44:46 +0000
Subject: Correctly bookkeep the size of TreeCache

---
 synapse/util/caches/treecache.py | 31 ++++++++++++++++++++++++++++---
 1 file changed, 28 insertions(+), 3 deletions(-)

(limited to 'synapse/util/caches/treecache.py')

diff --git a/synapse/util/caches/treecache.py b/synapse/util/caches/treecache.py
index a29ea8144e..3331ea9eba 100644
--- a/synapse/util/caches/treecache.py
+++ b/synapse/util/caches/treecache.py
@@ -21,7 +21,7 @@ class TreeCache(object):
         node = self.root
         for k in key[:-1]:
             node = node.setdefault(k, {})
-        node[key[-1]] = value
+        node[key[-1]] = _Entry(value)
         self.size += 1
 
     def get(self, key, default=None):
@@ -30,7 +30,7 @@ class TreeCache(object):
             node = node.get(k, None)
             if node is None:
                 return default
-        return node.get(key[-1], default)
+        return node.get(key[-1], _Entry(default)).value
 
     def clear(self):
         self.size = 0
@@ -60,8 +60,33 @@ class TreeCache(object):
                 break
             node_and_keys[i+1][0].pop(k)
 
-        self.size -= 1
+        popped, cnt = _strip_and_count_entires(popped)
+        self.size -= cnt
         return popped
 
     def __len__(self):
         return self.size
+
+
+class _Entry(object):
+    __slots__ = ["value"]
+
+    def __init__(self, value):
+        object.__setattr__(self, "value", value)
+
+
+def _strip_and_count_entires(d):
+    """Takes an _Entry or dict with leaves of _Entry's, and either returns the
+    value or a dictionary with _Entry's replaced by their values.
+
+    Also returns the count of _Entry's
+    """
+    if isinstance(d, dict):
+        cnt = 0
+        for key, value in d.items():
+            v, n = _strip_and_count_entires(value)
+            d[key] = v
+            cnt += n
+        return d, cnt
+    else:
+        return d.value, 1
-- 
cgit 1.4.1


From fb7299800feb0b8bc0f8429fc5b840abea0609d3 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Fri, 29 Jan 2016 11:29:14 +0000
Subject: Directly set self.value

---
 synapse/util/caches/treecache.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'synapse/util/caches/treecache.py')

diff --git a/synapse/util/caches/treecache.py b/synapse/util/caches/treecache.py
index 3331ea9eba..29d02f7e95 100644
--- a/synapse/util/caches/treecache.py
+++ b/synapse/util/caches/treecache.py
@@ -72,7 +72,7 @@ class _Entry(object):
     __slots__ = ["value"]
 
     def __init__(self, value):
-        object.__setattr__(self, "value", value)
+        self.value = value
 
 
 def _strip_and_count_entires(d):
-- 
cgit 1.4.1