Skip to content

Instantly share code, notes, and snippets.

@njsmith
Created Jul 31, 2020
Embed
What would you like to do?
import trio
import weakref
import cachetools
# Note: this assumes you only have one trio thread running at once...
# In the unlikely event that you have multiple, you should put the locks and cache
# into thread-local or run-local storage.
def trio_lru_cache(fn):
cache = cachetools.LRUCache()
# Trick: this contains a lock for each unique cache key in use,
# that we use to make sure we don't have multiple tasks trying
# to rebuild the same cache key at once. Since it's a WeakValueDictionary,
# unused keys are automatically evicted by the GC.
locks = weakref.WeakValueDictionary()
@functools.wraps(fn)
async def wrapper(*args, **kwargs):
key = make_key(args, kwargs)
try:
return cache[key]
except KeyError:
lock = locks.setdefault(key, trio.Lock())
async with lock:
try:
return cache[key]
except KeyError:
new_value = await fn(*args, **kwargs)
cache[key] = new_value
return new_value
return wrapper
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment