Source code for baseplate.lib.ratelimit.backends.memcache

from pymemcache.client import PooledClient

from baseplate import Span
from baseplate.clients import ContextFactory
from baseplate.clients.memcache import MemcacheContextFactory
from baseplate.clients.memcache import MonitoredMemcacheConnection
from baseplate.lib.ratelimit.backends import _get_current_bucket
from baseplate.lib.ratelimit.backends import RateLimitBackend


[docs]class MemcacheRateLimitBackendContextFactory(ContextFactory): """MemcacheRateLimitBackend context factory. :param memcache_pool: The memcache pool to back this ratelimit. :param prefix: A prefix to add to keys during rate limiting. This is useful if you will have two different rate limiters that will receive the same keys. """ def __init__(self, memcache_pool: PooledClient, prefix: str = "rl:"): self.memcache_context_factory = MemcacheContextFactory(memcache_pool) self.prefix = prefix
[docs] def make_object_for_context(self, name: str, span: Span) -> "MemcacheRateLimitBackend": memcache = self.memcache_context_factory.make_object_for_context(name, span) return MemcacheRateLimitBackend(memcache, prefix=self.prefix)
[docs]class MemcacheRateLimitBackend(RateLimitBackend): """A Memcache backend for rate limiting. :param memcache: A memcached connection. :param prefix: A prefix to add to keys during rate limiting. This is useful if you will have two different rate limiters that will receive the same keys. """ def __init__(self, memcache: MonitoredMemcacheConnection, prefix: str = "rl:"): self.memcache = memcache self.prefix = prefix
[docs] def consume(self, key: str, amount: int, allowance: int, interval: int) -> bool: """Consume the given `amount` from the allowance for the given `key`. This will return true if the `key` remains below the `allowance` after consuming the given `amount`. :param key: The name of the rate limit bucket to consume from. :param amount: The amount to consume from the rate limit bucket. :param allowance: The maximum allowance for the rate limit bucket. :param interval: The interval to reset the allowance. """ current_bucket = _get_current_bucket(interval) key = self.prefix + key + current_bucket ttl = interval * 2 self.memcache.add(key, 0, expire=ttl) # `incr` will return None if we experience a delay after the prior # `add` call that causes the ttl to expire. We default to `amount` in # this case. count = self.memcache.incr(key, amount) or amount return count <= allowance