Something went wrong. Try again.
🌉 A bridge between decentralized social networks fed.brid.gy
Something went wrong. Try again.
bridgy-fed memcache.py
12 kB · 358 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359"""Utilities for caching data in memcache.
TODO: move most or all of this to webutil?"""from datetime import datetime, timedelta, timezonefrom enum import auto, Enumimport functoolsimport inspectimport loggingimport osimport reimport string
import configfrom google.cloud.ndb._cache import global_cache_keyfrom google.cloud.ndb.global_cache import _InProcessGlobalCache, MemcacheCachefrom google.cloud.ndb.key import Keyfrom pymemcache.client.base import PooledClientfrom pymemcache.serde import PickleSerdefrom pymemcache.test.utils import MockMemcacheClientfrom webutil import appengine_info, util
from domains import PRIMARY_DOMAIN
logger = logging.getLogger(__name__)
# https://github.com/memcached/memcached/wiki/Commands#standard-protocolKEY_MAX_LEN = 250
MEMOIZE_VERSION = 2
# per-user rates for running tasks. rate limits and spreads out tasks for bursty# users. values map protocol label to delay. None means all protocols.# https://github.com/snarfed/bridgy-fed/issues/1788PER_USER_TASK_RATES = { 'receive': { None: timedelta(seconds=5), # all protocols }, 'send': { 'atproto': timedelta(seconds=10), },}
class RateLimitType(Enum): LINEAR = auto() EXPONENTIAL = auto()
NDB_MEMCACHE_TIMEOUT = timedelta(hours=2)
WHITESPACE_RE = re.compile(f'[{string.whitespace}]')
# https://pymemcache.readthedocs.io/en/latest/apidoc/pymemcache.client.base.html#pymemcache.client.base.Client.__init__kwargs = { 'server': os.environ.get('MEMCACHE_HOST', 'localhost'), 'allow_unicode_keys': True, 'default_noreply': False, 'timeout': 10, # seconds 'connect_timeout': 10, # seconds}
if appengine_info.DEBUG or appengine_info.LOCAL_SERVER: logger.info(f'Using in memory mock memcache: {kwargs}') memcache = MockMemcacheClient(**kwargs) pickle_memcache = MockMemcacheClient(serde=PickleSerde(), **kwargs) global_cache = _InProcessGlobalCache()else: logger.info(f'Using production Memorystore memcache: {kwargs}') memcache = PooledClient(**kwargs) pickle_memcache = PooledClient(serde=PickleSerde(), **kwargs) global_cache = MemcacheCache(memcache, strict_read=False, strict_write=False)
def cache_policy(key): """In memory ndb cache.
https://github.com/snarfed/bridgy-fed/issues/1149#issuecomment-2261383697
Only cache kinds in memory that are immutable or largely harmless when changed.
Keep an eye on this in case we start seeing problems due to this ndb bug where unstored in-memory modifications get returned by later gets: https://github.com/googleapis/python-ndb/issues/888
Args: key (google.cloud.datastore.key.Key or google.cloud.ndb.key.Key): see https://github.com/googleapis/python-ndb/issues/987
Returns: bool: whether to cache this object """ if isinstance(key, Key): # use internal google.cloud.datastore.key.Key # https://github.com/googleapis/python-ndb/issues/987 key = key._key
return key and key.kind in ('AtpBlock', 'Object')
def global_cache_policy(key): return True
def global_cache_timeout_policy(key): """Cache everything for 2h.
Args: key (google.cloud.datastore.key.Key or google.cloud.ndb.key.Key): see https://github.com/googleapis/python-ndb/issues/987
Returns: int: cache expiration for this object, in seconds """ if isinstance(key, Key): # use internal google.cloud.datastore.key.Key # https://github.com/googleapis/python-ndb/issues/987 key = key._key
if key.kind == 'AtpBlock': return None
return int(NDB_MEMCACHE_TIMEOUT.total_seconds())
def key(key): """Preprocesses a memcache key. Right now just truncates it to 250 chars.
https://pymemcache.readthedocs.io/en/latest/apidoc/pymemcache.client.base.html https://github.com/memcached/memcached/wiki/Commands#standard-protocol
TODO: truncate to 250 *UTF-8* chars, to handle Unicode chars in URLs. Related: pymemcache Client's allow_unicode_keys constructor kwarg.
Args: key (str)
Returns: bytes: """ assert isinstance(key, str), repr(key) return WHITESPACE_RE.sub('_', key).encode()[:KEY_MAX_LEN]
def memoize_key(fn, *args, _version=MEMOIZE_VERSION, **kwargs): # normalize kwargs with defaults so that eg fn(x) and fn(x, y=z) get # the same key when y defaults to z try: bound = inspect.signature(fn).bind(*args, **kwargs) bound.apply_defaults() args, kwargs = bound.args, bound.kwargs except TypeError: # not fn's own args, eg memoize's key kwarg passes a synthetic key pass
return key(f'{fn.__qualname__}-{_version}-{repr(args)}-{repr(kwargs)}')
NONE = () # empty tuple
def memoize(expire=None, key=None, write=True, version=MEMOIZE_VERSION): """Memoize function decorator that stores the cached value in memcache.
Args: expire (datetime.timedelta): optional, expiration key (callable): function that takes the function's ``(*args, **kwargs)`` and returns the cache key to use. If it returns None, memcache won't be used. write (bool or callable): whether to write to memcache. If this is a callable, it will be called with the function's ``(*args, **kwargs)`` and should return True or False. version (int): overrides our default version number in the memcache key. Bumping this version can have the same effect as clearing the cache for just the affected function. """ expire = int(expire.total_seconds()) if expire else 0
def decorator(fn): @functools.wraps(fn) def wrapped(*args, **kwargs): cache_key = None if key: key_val = key(*args, **kwargs) if key_val: cache_key = memoize_key(fn, key_val, _version=version) else: cache_key = memoize_key(fn, *args, _version=version, **kwargs)
if pickle_memcache and cache_key: val = pickle_memcache.get(cache_key) if val is not None: logger.debug(f'cache hit {cache_key} {repr(val)[:100]}') return None if val == NONE else val else: logger.debug(f'cache miss {cache_key}')
val = fn(*args, **kwargs)
if pickle_memcache and cache_key: write_cache = (write if isinstance(write, bool) else write(*args, **kwargs)) if write_cache: logger.debug(f'cache set {cache_key} {repr(val)[:100]}') pickle_memcache.set(cache_key, NONE if val is None else val, expire=expire)
return val
return wrapped
return decorator
def evict(entity_key): """Evict a datastore entity from memcache.
For :class:`models.User` and :class:`models.Object` entities, also clears their copies from the :func:`models.get_original_user_key` and :func:`models.get_original_object_key` memoize caches.
Args: entity_key (google.cloud.ndb.Key) """ if entity := entity_key.get(): for val in getattr(entity, 'copies', []): entity.clear_get_original_cache(val.uri)
global_cache.delete([global_cache_key(entity_key._key)])
def evict_raw(key): """Evict a key from memcache.
Args: key (str)
Returns: bool: whether the key existed and was deleted """ return memcache.delete(key)
def remote_evict(entity_key=None, raw=None): """Send a request to production Bridgy Fed to evict a key from memcache.
Args: entity_key (google.cloud.ndb.Key) raw (str)
Returns: requests.Response: """ if entity_key: data = {'key': entity_key.urlsafe()} elif raw: data = {'raw': raw} else: raise ValueError('either entity_key or raw are required')
return util.requests_post(f'https://{PRIMARY_DOMAIN}/admin/memcache/evict', headers={'Authorization': config.SECRET_KEY}, data=data)
def task_eta(queue, user_id, protocol=None): """Get the ETA to use for a given user's task in a given queue.
Task rate limit delays are per user, stored in memcache with a key based on ``queue`` and ``user_id`` and an integer value of POSIX timestamp (UTC) in seconds.
Only generates ETAs for task queues in :attr:`PER_USER_TASK_RATES`. Calls for other queues always return ``None``.
Background: https://github.com/snarfed/bridgy-fed/issues/1788
Args: queue (str) user_id (str) protocol (str): optional protocol label to look up protocol-specific delay and :class:`RateLimitType`
Returns: datetime.datetime: the ETA for this task, or ``None`` if the ETA is now """ from models import PROTOCOLS
if not (delays := PER_USER_TASK_RATES.get(queue)): return None
# look up delay for protocol, fall back to None (all protocols) if not (delay := delays.get(protocol) or delays.get(None)): return None
cache_key = key(f'task-delay-{queue}-{user_id}')
now = util.now()
if protocol and PROTOCOLS[protocol].RATE_LIMIT_TYPE == RateLimitType.EXPONENTIAL: if eta_s := memcache.get(cache_key): eta = datetime.fromtimestamp(eta_s, timezone.utc) if eta >= now: cur_delay = eta - now new_eta = eta + max(cur_delay, delay) memcache.set(cache_key, int(new_eta.timestamp())) return new_eta
else: # linear if eta_s := memcache.incr(cache_key, int(delay.total_seconds())): eta = datetime.fromtimestamp(eta_s, timezone.utc) if eta > now: return eta
# incr failed (key doesn't exist) or timestamp is in the past, set it to now # # note that this isn't synchronized; multiple callers may race and both get now # as the returned ETA. that's ok, we don't depend on this for correctness in any # way, just best-effort rate limiting. memcache.set(cache_key, int(now.timestamp())) return now
###########################################
# https://github.com/googleapis/python-ndb/issues/743#issuecomment-2067590945## fixes "RuntimeError: Key has already been set in this batch" errors due to# tasklets in pages.serve_feedfrom logging import error as log_errorfrom sys import modules
from google.cloud.ndb._cache import ( _GlobalCacheSetBatch, global_compare_and_swap, global_set_if_not_exists, global_watch,)from google.cloud.ndb.tasklets import Future, Return, tasklet
GLOBAL_CACHE_KEY_PREFIX: bytes = modules["google.cloud.ndb._cache"]._PREFIXLOCKED_FOR_READ: bytes = modules["google.cloud.ndb._cache"]._LOCKED_FOR_READLOCK_TIME: bytes = modules["google.cloud.ndb._cache"]._LOCK_TIME
@taskletdef custom_global_lock_for_read(key: str, value: str): if value is not None: yield global_watch(key, value) lock_acquired = yield global_compare_and_swap( key, LOCKED_FOR_READ, expires=LOCK_TIME ) else: lock_acquired = yield global_set_if_not_exists( key, LOCKED_FOR_READ, expires=LOCK_TIME )
if lock_acquired: raise Return(LOCKED_FOR_READ)
modules["google.cloud.ndb._cache"].global_lock_for_read = custom_global_lock_for_read