Rate Limite

lua_script.py

# Leaky bucket implementation from:
# https://blog.callr.tech/rate-limiting-for-distributed-systems-with-redis-and-lua/

LEAKY_BUCKET_SCRIPT = b"""
    local ts  = tonumber(ARGV[1]) -- current timestamp
    local cps = tonumber(ARGV[2]) -- calls per second
    local key = KEYS[1]
    
    -- remove tokens older than 1 second ago
    local min = ts - 1
    redis.call('ZREMRANGEBYSCORE', key, '-inf', min)

    -- get the last item of the sorted set
    local last = redis.call('ZRANGE', key, -1, -1)

    -- the default value of next if there are no previous items is the current timestamp
    local next = ts
    
    if type(last) == 'table' and #last > 0 then
      -- this iteration only serves to access the value in the set
      
      for key, value in pairs(last) do
        -- the next value is the last timestamp stored plus
        -- x miliseconds depending on the allowed cps
        next = tonumber(value) + 1/cps
        
        break -- break at first item
      end
    end

    -- if the current ts is > than last+1/cps we'll use ts
    if ts > next then
      next = ts
    end

    -- store the new timestamp in the set
    redis.call('ZADD', key, next, next)
    
    -- return how long it should wait in order to respect the rate limit
    return tostring(next - ts)
"""

LEAKY_BUCKET_SCRIPT_HASH = sha1(LEAKY_BUCKET_SCRIPT).hexdigest()  # nosec

rate_limit.py

# Code adapted from:
# https://github.com/EvoluxBR/python-redis-rate-limit

import time
from distutils.version import StrictVersion
from hashlib import sha1

from django.conf import settings
from django.utils import timezone

import redis
from redis.exceptions import NoScriptError


class RedisVersionNotSupported(Exception):
    """
    Rate Limit depends on Redis' commands EVALSHA and EVAL which are
    only available since the version 2.6.0 of the database.
    """


class RateLimit:
    """
    This class offers an abstraction of a Rate Limit algorithm implemented on
    top of Redis >= 2.6.0.
    """

    def __init__(self, resource, max_requests):
        """
        Class initialization method checks if the Rate Limit algorithm is
        actually supported by the installed Redis version and sets some
        useful properties.
        If Rate Limit is not supported, it raises an Exception.
        :param resource: resource identifier string (i.e. 'user_pictures')
        :param max_requests: integer (i.e. '10')
        """
        self._client = redis.Redis.from_url(settings.REDIS_URL)
        if not self._is_rate_limit_supported():
            raise RedisVersionNotSupported()

        self._rate_limit_key = "rate_limit:{0}".format(resource)
        self._max_requests = max_requests

    def __enter__(self):
        wait_for = self.get_rate_limit_wait()
        time.sleep(wait_for)

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def get_rate_limit_wait(self):
        """
        Calls a LUA script that should increment the resource usage by client.
        If the resource limit overflows the maximum number of requests, this
        method raises an Exception.
        :return: integer: current usage
        """
        ts = timezone.now().timestamp()

        try:
            wait_for = self._client.evalsha(
                LEAKY_BUCKET_SCRIPT_HASH, 1, self._rate_limit_key, ts, self._max_requests
            )
        except NoScriptError:
            wait_for = self._client.eval(
                LEAKY_BUCKET_SCRIPT, 1, self._rate_limit_key, ts, self._max_requests
            )

        return float(wait_for)

    def _is_rate_limit_supported(self):
        """
        Checks if Rate Limit is supported which can basically be found by
        looking at Redis database version that should be 2.6.0 or greater.
        :return: bool
        """
        redis_version = self._client.info()["redis_version"]
        is_supported = StrictVersion(redis_version) >= StrictVersion("2.6.0")
        return bool(is_supported)

    def _reset(self):
        """
        Deletes all keys that start with 'rate_limit:'.
        """
        for rate_limit_key in self._client.keys("rate_limit:*"):
            self._client.delete(rate_limit_key)

simplified.py

import time
from hashlib import sha1

from django.utils import timezone

import redis
from redis.exceptions import NoScriptError


def sleep_to_rate_limit(key, requests_per_minute):
    client = redis.Redis.from_url(settings.REDIS_URL)
    rate_limit_key = "rate_limit:{0}".format(key)
    requests_per_second = requests_per_minute / 60
    current_time = timezone.now().timestamp()

    try:
        wait_for = client.evalsha(
            LEAKY_BUCKET_SCRIPT_HASH, 1, rate_limit_key, current_time, requests_per_second
        )
    except NoScriptError:
        wait_for = client.eval(
            LEAKY_BUCKET_SCRIPT, 1, rate_limit_key, current_time, requests_per_second
        )

    time.sleep(float(wait_for))