From 46aebbbcbf94eb78ae45d3bb3bf3ffeabb44dd4f Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Mon, 16 Jan 2017 13:48:04 +0000
Subject: [PATCH] Add support for 'iterable' to ExpiringCache

---
 synapse/state.py                     |  6 +++++-
 synapse/util/caches/expiringcache.py | 26 +++++++++++++++++---------
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/synapse/state.py b/synapse/state.py
index b9d5627a82..461e82acdf 100644
--- a/synapse/state.py
+++ b/synapse/state.py
@@ -41,7 +41,7 @@ KeyStateTuple = namedtuple("KeyStateTuple", ("context", "type", "state_key"))
 CACHE_SIZE_FACTOR = float(os.environ.get("SYNAPSE_CACHE_FACTOR", 0.1))
 
 
-SIZE_OF_CACHE = int(1000 * CACHE_SIZE_FACTOR)
+SIZE_OF_CACHE = int(10000 * CACHE_SIZE_FACTOR)
 EVICTION_TIMEOUT_SECONDS = 60 * 60
 
 
@@ -77,6 +77,9 @@ class _StateCacheEntry(object):
         else:
             self.state_id = _gen_state_id()
 
+    def __len__(self):
+        return len(self.state)
+
 
 class StateHandler(object):
     """ Responsible for doing state conflict resolution.
@@ -99,6 +102,7 @@ class StateHandler(object):
             clock=self.clock,
             max_len=SIZE_OF_CACHE,
             expiry_ms=EVICTION_TIMEOUT_SECONDS * 1000,
+            iterable=True,
             reset_expiry_on_get=True,
         )
 
diff --git a/synapse/util/caches/expiringcache.py b/synapse/util/caches/expiringcache.py
index 080388958f..9b44b3fab3 100644
--- a/synapse/util/caches/expiringcache.py
+++ b/synapse/util/caches/expiringcache.py
@@ -23,7 +23,7 @@ logger = logging.getLogger(__name__)
 
 class ExpiringCache(object):
     def __init__(self, cache_name, clock, max_len=0, expiry_ms=0,
-                 reset_expiry_on_get=False):
+                 reset_expiry_on_get=False, iterable=False):
         """
         Args:
             cache_name (str): Name of this cache, used for logging.
@@ -36,6 +36,8 @@ class ExpiringCache(object):
                 evicted based on time.
             reset_expiry_on_get (bool): If true, will reset the expiry time for
                 an item on access. Defaults to False.
+            iterable (bool): If true, the size is calculated by summing the
+                sizes of all entries, rather than the number of entries.
 
         """
         self._cache_name = cache_name
@@ -49,7 +51,9 @@ class ExpiringCache(object):
 
         self._cache = {}
 
-        self.metrics = register_cache(cache_name, self._cache)
+        self.metrics = register_cache(cache_name, self)
+
+        self.iterable = iterable
 
     def start(self):
         if not self._expiry_ms:
@@ -66,14 +70,15 @@ class ExpiringCache(object):
         self._cache[key] = _CacheEntry(now, value)
 
         # Evict if there are now too many items
-        if self._max_len and len(self._cache.keys()) > self._max_len:
+        if self._max_len and len(self) > self._max_len:
             sorted_entries = sorted(
-                self._cache.items(),
+                self._cache.keys(),
                 key=lambda item: item[1].time,
             )
 
-            for k, _ in sorted_entries[self._max_len:]:
-                self._cache.pop(k)
+            while len(self) > self._max_len and sorted_entries:
+                key = sorted_entries.pop()
+                self._cache.pop(key)
 
     def __getitem__(self, key):
         try:
@@ -99,7 +104,7 @@ class ExpiringCache(object):
             # zero expiry time means don't expire. This should never get called
             # since we have this check in start too.
             return
-        begin_length = len(self._cache)
+        begin_length = len(self)
 
         now = self._clock.time_msec()
 
@@ -114,11 +119,14 @@ class ExpiringCache(object):
 
         logger.debug(
             "[%s] _prune_cache before: %d, after len: %d",
-            self._cache_name, begin_length, len(self._cache)
+            self._cache_name, begin_length, len(self)
         )
 
     def __len__(self):
-        return len(self._cache)
+        if self.iterable:
+            return sum(len(value.value) for value in self._cache.itervalues())
+        else:
+            return len(self._cache)
 
 
 class _CacheEntry(object):