diff --git a/synapse/util/caches/lrucache.py b/synapse/util/caches/lrucache.py
index 9c4c679175..cf5fbb679c 100644
--- a/synapse/util/caches/lrucache.py
+++ b/synapse/util/caches/lrucache.py
@@ -49,7 +49,7 @@ class LruCache(object):
Can also set callbacks on objects when getting/setting which are fired
when that key gets invalidated/evicted.
"""
- def __init__(self, max_size, keylen=1, cache_type=dict):
+ def __init__(self, max_size, keylen=1, cache_type=dict, size_callback=None):
cache = cache_type()
self.cache = cache # Used for introspection.
list_root = _Node(None, None, None, None)
@@ -58,6 +58,12 @@ class LruCache(object):
lock = threading.Lock()
+ def evict():
+ while cache_len() > max_size:
+ todelete = list_root.prev_node
+ delete_node(todelete)
+ cache.pop(todelete.key, None)
+
def synchronized(f):
@wraps(f)
def inner(*args, **kwargs):
@@ -66,6 +72,16 @@ class LruCache(object):
return inner
+ cached_cache_len = [0]
+ if size_callback is not None:
+ def cache_len():
+ return cached_cache_len[0]
+ else:
+ def cache_len():
+ return len(cache)
+
+ self.len = synchronized(cache_len)
+
def add_node(key, value, callbacks=set()):
prev_node = list_root
next_node = prev_node.next_node
@@ -74,6 +90,9 @@ class LruCache(object):
next_node.prev_node = node
cache[key] = node
+ if size_callback:
+ cached_cache_len[0] += size_callback(node.value)
+
def move_node_to_front(node):
prev_node = node.prev_node
next_node = node.next_node
@@ -92,23 +111,25 @@ class LruCache(object):
prev_node.next_node = next_node
next_node.prev_node = prev_node
+ if size_callback:
+ cached_cache_len[0] -= size_callback(node.value)
+
for cb in node.callbacks:
cb()
node.callbacks.clear()
@synchronized
- def cache_get(key, default=None, callback=None):
+ def cache_get(key, default=None, callbacks=[]):
node = cache.get(key, None)
if node is not None:
move_node_to_front(node)
- if callback:
- node.callbacks.add(callback)
+ node.callbacks.update(callbacks)
return node.value
else:
return default
@synchronized
- def cache_set(key, value, callback=None):
+ def cache_set(key, value, callbacks=[]):
node = cache.get(key, None)
if node is not None:
if value != node.value:
@@ -116,21 +137,18 @@ class LruCache(object):
cb()
node.callbacks.clear()
- if callback:
- node.callbacks.add(callback)
+ if size_callback:
+ cached_cache_len[0] -= size_callback(node.value)
+ cached_cache_len[0] += size_callback(value)
+
+ node.callbacks.update(callbacks)
move_node_to_front(node)
node.value = value
else:
- if callback:
- callbacks = set([callback])
- else:
- callbacks = set()
- add_node(key, value, callbacks)
- if len(cache) > max_size:
- todelete = list_root.prev_node
- delete_node(todelete)
- cache.pop(todelete.key, None)
+ add_node(key, value, set(callbacks))
+
+ evict()
@synchronized
def cache_set_default(key, value):
@@ -139,10 +157,7 @@ class LruCache(object):
return node.value
else:
add_node(key, value)
- if len(cache) > max_size:
- todelete = list_root.prev_node
- delete_node(todelete)
- cache.pop(todelete.key, None)
+ evict()
return value
@synchronized
@@ -174,10 +189,8 @@ class LruCache(object):
for cb in node.callbacks:
cb()
cache.clear()
-
- @synchronized
- def cache_len():
- return len(cache)
+ if size_callback:
+ cached_cache_len[0] = 0
@synchronized
def cache_contains(key):
@@ -190,7 +203,7 @@ class LruCache(object):
self.pop = cache_pop
if cache_type is TreeCache:
self.del_multi = cache_del_multi
- self.len = cache_len
+ self.len = synchronized(cache_len)
self.contains = cache_contains
self.clear = cache_clear
|
