Implement a throttle class

[jkemp101]

A simple throttle class might be a better interface to the existing throttle functions. The global redis instance limits the ability to use multiple instances easily and generally isn't a great pattern. The alternative of passing a redis instance to each function doesn't seem great either. A class would make it easily to encapsulate some of this logic and settings.

[froxCZ]

@jkemp101

Spec

Below are things worth mentioning as a result of my research and POC.

Throttle Class API

• Almost every method should accept suffix, to not require specific class instance when a resource is per entity so it can be used e.g. for api rate limit company like this Throttle('api_access').throttle(suffix=company_id)
• Constructor should accept redis instance, name of resource and default config values (knobs).
• set_knobs should set the knobs to redis, same as throttle_set did.
• Throttle lua script: The script is tied to redis instance, so ideally, we should register the lua script only once per redis instance. We can either:
  • use a global dict where key is a redis instance and value is the lua script. To still allow the redis instance to be garbage collected if not used anymore, we should use WeakKeyDictionary.
  • Abuse the redis class (which we don't own) and store the throttle script on it (redis_instance.throttle_script = throttle_script). Next time, do not register a script if given redis instance already has that property.
  • Introduce our LimitLionRedisInstance class, with reference to redis and throttle_script. Every redis instance would be wrapped in this class, and the reference to it would be managed by caller code. Throttle would accept LimitLionRedisInstance instead of Redis instance in constructor.
• Moving other throttle_* functions to class should be straightforward
• Bonus:def throttle_configure(redis_instance, testing=False): fn has the testing flag to fix time in the lua script for testing purposes. Ideally, we could use some mock or patch and remove code related to testing from production code.

Constructor

Default Throttle configs should be passed to the class constructor, together with redis instance and the throttle name

def __init__(self,
             redis_instance,
             name,
             default_rps,
             default_burst=THROTTLE_BURST_DEFAULT,
             default_window=THROTTLE_WINDOW_DEFAULT,
             default_knobs_ttl=DEFAULT_KNOBS_TTL,
             testing=False):

Changing the config in runtime

def set_knobs(self, suffix=None, rps=None, burst=None, window=None,
              knobs_ttl=None):

POC

Example of Throttle class storing redis and lua script via WeakKeyDictionary and example of throttle fn

import time
from weakref import WeakKeyDictionary

import pkg_resources

KEY_FORMAT = 'throttle:{}'

# throttle knob defaults
THROTTLE_BURST_DEFAULT = 1
THROTTLE_WINDOW_DEFAULT = 5
THROTTLE_REQUESTED_TOKENS_DEFAULT = 1

# The default is to extend a throttle's knob settings TTL out
# 7 days each time the throttle is used.
DEFAULT_KNOBS_TTL = 60 * 60 * 24 * 7


REDIS_TO_SCRIPT_CACHE = WeakKeyDictionary()

class Throttle:
    def __init__(self,
                 redis_instance, name, default_rps,
                 default_burst=THROTTLE_BURST_DEFAULT,
                 default_window=THROTTLE_WINDOW_DEFAULT,
                 default_knobs_ttl=DEFAULT_KNOBS_TTL,
                 testing=False):
        self.redis_instance = redis_instance
        self.name = name
        self.default_rps = default_rps
        self.default_burst = default_burst
        self.default_window = default_window
        self.default_knobs_ttl = default_knobs_ttl
        self.throttle_script = self._register_throttle_script(testing)

    def _register_throttle_script(self, testing=False):
        script = REDIS_TO_SCRIPT_CACHE.get(self.redis_instance)
        if script:
            return script
        lua_script = pkg_resources.resource_string(
            __name__, 'throttle.lua'
        ).decode()

        # Modify scripts when testing so time can be frozen
        if testing:
            lua_script = lua_script.replace(
                'local time = redis.call("time")',
                'local time\n'
                'if redis.call("exists", "frozen_second") == 1 then\n'
                '  time = redis.call("mget", "frozen_second", "frozen_microsecond")\n'  # noqa: E501
                'else\n'
                '  time = redis.call("time")\n'
                'end',
            )
        throttle_script = self.redis_instance.register_script(lua_script)
        REDIS_TO_SCRIPT_CACHE[self.redis_instance] = throttle_script
        return throttle_script

    def _key(self, suffix=None):
        full_name = self.name
        if suffix:
            full_name += suffix
        return KEY_FORMAT.format(full_name)

    def throttle(self, requested_tokens=THROTTLE_REQUESTED_TOKENS_DEFAULT, suffix=None):
        allowed, tokens, sleep = self.throttle_script(
            keys=[],
            args=[
                self._key(suffix),
                self.default_rps,
                self.default_burst,
                self.default_window,
                self.requested_tokens,
                self.default_knobs_ttl,
            ],
        )
        # Converting the string sleep to a float causes floating point rounding
        # issues that limits having true microsecond resolution for the sleep
        # value.
        return allowed == 1, int(tokens), float(sleep)

    def set_knobs(self, suffix=None, rps=None, burst=None, window=None,
                  knobs_ttl=None):
        # update knobs in only in redis
        pass