From 1d894e1113ab834fd84de3dc3adebff1950b8450 Mon Sep 17 00:00:00 2001 From: Andrej Romanov Date: Thu, 19 Jun 2025 21:43:02 +0300 Subject: [PATCH 1/6] use redis as a storage for throttling statistic --- pyproject.toml | 5 +- tests/apps.py | 27 +++++++- tests/{ => web}/test_throttling.py | 0 winter/web/__init__.py | 6 +- winter/web/throttling/__init__.py | 6 ++ .../redis_throttling_configuration.py | 9 +++ winter/web/{ => throttling}/throttling.py | 28 +++++---- .../throttling_statistic_storage.py | 63 +++++++++++++++++++ 8 files changed, 129 insertions(+), 15 deletions(-) rename tests/{ => web}/test_throttling.py (100%) create mode 100644 winter/web/throttling/__init__.py create mode 100644 winter/web/throttling/redis_throttling_configuration.py rename winter/web/{ => throttling}/throttling.py (76%) create mode 100644 winter/web/throttling/throttling_statistic_storage.py diff --git a/pyproject.toml b/pyproject.toml index e29423b1..95203a41 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "winter" -version = "30.0.1" +version = "31.0.0" homepage = "https://github.com/WinterFramework/winter" description = "Web Framework with focus on python typing, dataclasses and modular design" authors = ["Alexander Egorov "] @@ -41,6 +41,7 @@ pydantic = ">=1.10, <2" openapi-spec-validator = ">=0.5.7, <1" uritemplate = "==4.2.0" # Lib doesn't follow semantic versioning httpx = ">=0.24.1, <0.28" +redis = "^6.2.0" [tool.poetry.dev-dependencies] flake8 = ">=3.7.7, <4" @@ -61,6 +62,8 @@ pytz = ">=2020.5" [tool.poetry.group.dev.dependencies] setuptools = "^71.1.0" +testcontainers = "^4.10.0" +redis = "^6.2.0" [build-system] requires = ["poetry-core>=1.3.1"] diff --git a/tests/apps.py b/tests/apps.py index 7c205b22..434753a5 100644 --- a/tests/apps.py +++ b/tests/apps.py @@ -1,6 +1,10 @@ +import atexit + from django.apps import AppConfig +from testcontainers.redis import RedisContainer from tests.web.interceptors import HelloWorldInterceptor +from winter.web import RedisThrottlingConfiguration from winter.web import exception_handlers_registry from winter.web import interceptor_registry from winter.web.exceptions.handlers import DefaultExceptionHandler @@ -9,6 +13,10 @@ class TestAppConfig(AppConfig): name = 'tests' + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._redis_container: RedisContainer | None = None + def ready(self): # define this import for force initialization all modules and to register Exceptions from .urls import urlpatterns # noqa: F401 @@ -19,7 +27,24 @@ def ready(self): interceptor_registry.add_interceptor(HelloWorldInterceptor()) winter_openapi.setup() - winter.web.setup() + + self._redis_container = RedisContainer() + self._redis_container.start() + self._redis_container.get_client().flushdb() + atexit.register(self.cleanup_redis) + + redis_throttling_configuration = RedisThrottlingConfiguration( + host=self._redis_container.get_container_host_ip(), + port=self._redis_container.port, + db=0, + password=self._redis_container.password + ) + winter.web.setup(redis_throttling_configuration) + winter_django.setup() exception_handlers_registry.set_default_handler(DefaultExceptionHandler) # for 100% test coverage + + def cleanup_redis(self): + if self._redis_container: + self._redis_container.stop() diff --git a/tests/test_throttling.py b/tests/web/test_throttling.py similarity index 100% rename from tests/test_throttling.py rename to tests/web/test_throttling.py diff --git a/winter/web/__init__.py b/winter/web/__init__.py index 4372de0f..16c467d1 100644 --- a/winter/web/__init__.py +++ b/winter/web/__init__.py @@ -20,11 +20,13 @@ from .response_header_resolver import ResponseHeaderArgumentResolver from .response_header_serializer import response_headers_serializer from .response_status_annotation import response_status +from .throttling import init_throttling_statistic_storage +from .throttling import RedisThrottlingConfiguration from .throttling import throttling from .urls import register_url_regexp -def setup(): +def setup(redis_throttling_configuration: RedisThrottlingConfiguration): from winter.data.exceptions import NotFoundException from .exceptions import RedirectException from .exceptions.problem_handling import autodiscover_problem_annotations @@ -61,3 +63,5 @@ def setup(): } for exception_class, handler in auto_handle_exceptions.items(): exception_handlers_registry.add_handler(exception_class, handler, auto_handle=True) + + init_throttling_statistic_storage(redis_throttling_configuration) diff --git a/winter/web/throttling/__init__.py b/winter/web/throttling/__init__.py new file mode 100644 index 00000000..fe1f47d8 --- /dev/null +++ b/winter/web/throttling/__init__.py @@ -0,0 +1,6 @@ +from .throttling import throttling +from .throttling import reset +from .throttling import create_throttle_class +from .throttling_statistic_storage import init_throttling_statistic_storage +from .throttling_statistic_storage import get_throttling_statistic_storage +from .redis_throttling_configuration import RedisThrottlingConfiguration diff --git a/winter/web/throttling/redis_throttling_configuration.py b/winter/web/throttling/redis_throttling_configuration.py new file mode 100644 index 00000000..5f64b41d --- /dev/null +++ b/winter/web/throttling/redis_throttling_configuration.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass + + +@dataclass +class RedisThrottlingConfiguration: + host: str + port: int + db: int + password: str | None = None diff --git a/winter/web/throttling.py b/winter/web/throttling/throttling.py similarity index 76% rename from winter/web/throttling.py rename to winter/web/throttling/throttling.py index caf051e3..877963ea 100644 --- a/winter/web/throttling.py +++ b/winter/web/throttling/throttling.py @@ -5,12 +5,12 @@ from typing import Tuple import django.http -from django.core.cache import cache as default_cache from winter.core import annotate_method +from .throttling_statistic_storage import get_throttling_statistic_storage if TYPE_CHECKING: - from .routing import Route # noqa: F401 + from winter.web.routing import Route # noqa: F401 @dataclasses.dataclass @@ -33,23 +33,25 @@ def throttling(rate: Optional[str], scope: Optional[str] = None): class BaseRateThrottle: def __init__(self, throttling_: Throttling): self._throttling = throttling_ + self._storage = get_throttling_statistic_storage() def allow_request(self, request: django.http.HttpRequest) -> bool: ident = _get_ident(request) key = _get_cache_key(self._throttling.scope, ident) - history = default_cache.get(key, []) - now = time.time() + with self._storage.lock(key, timeout_sec=2): + history = self._storage.get(key) + now = time.time() - while history and history[-1] <= now - self._throttling.duration: - history.pop() + while history and history[-1] <= now - self._throttling.duration: + history.pop() - if len(history) >= self._throttling.num_requests: - return False + if len(history) >= self._throttling.num_requests: + return False - history.insert(0, now) - default_cache.set(key, history, self._throttling.duration) - return True + history.insert(0, now) + self._storage.set(key, history, self._throttling.duration) + return True def reset(request: django.http.HttpRequest, scope: str): @@ -59,7 +61,9 @@ def reset(request: django.http.HttpRequest, scope: str): """ ident = _get_ident(request) key = _get_cache_key(scope, ident) - default_cache.delete(key) + storage = get_throttling_statistic_storage() + with storage.lock(key, timeout_sec=2): + storage.delete(key) CACHE_KEY_FORMAT = 'throttle_{scope}_{ident}' diff --git a/winter/web/throttling/throttling_statistic_storage.py b/winter/web/throttling/throttling_statistic_storage.py new file mode 100644 index 00000000..b86d61be --- /dev/null +++ b/winter/web/throttling/throttling_statistic_storage.py @@ -0,0 +1,63 @@ +from contextlib import contextmanager +from typing import Generator + +from redis import Redis +from redis.exceptions import LockError +from redis.exceptions import LockNotOwnedError + +from .redis_throttling_configuration import RedisThrottlingConfiguration + + +class ThrottlingStatisticStorage: + def __init__(self, configuration: RedisThrottlingConfiguration): + self._redis_client = Redis( + host=configuration.host, + port=configuration.port, + db=configuration.db, + password=configuration.password, + decode_responses=True, + ) + + @contextmanager + def lock(self, key: str, timeout_sec: int) -> Generator[None, None, None]: + lock = self._redis_client.lock(f'lock:{key}', timeout=timeout_sec, blocking=True) + lock.acquire() + + try: + yield + finally: + try: + lock.release() + except (LockError, LockNotOwnedError): + pass + + def get(self, key: str) -> list[float]: + raw_value: str | None = self._redis_client.get(key) + if raw_value is None: + return [] + + value = map(float, raw_value.split(',')) + return list(value) + + def set(self, key: str, value: list[float], ex: int): + raw_value = ','.join(map(str, value)) + self._redis_client.set(key, raw_value, ex) + + def delete(self, key: str): + self._redis_client.delete(key) + + +_throttling_statistic_storage: ThrottlingStatisticStorage | None = None + + +def init_throttling_statistic_storage(configuration: RedisThrottlingConfiguration): + global _throttling_statistic_storage + if _throttling_statistic_storage is not None: + raise RuntimeError("ThrottlingStatisticStorage is already initialized.") + _throttling_statistic_storage = ThrottlingStatisticStorage(configuration) + + +def get_throttling_statistic_storage() -> ThrottlingStatisticStorage: + if _throttling_statistic_storage is None: + raise RuntimeError("ThrottlingStatisticStorage is not initialized.") + return _throttling_statistic_storage \ No newline at end of file From f8c8da1d730a889245e29b42673b32e9e6d58c33 Mon Sep 17 00:00:00 2001 From: Andrej Romanov Date: Mon, 23 Jun 2025 11:32:07 +0300 Subject: [PATCH 2/6] get rid of redis lock --- tests/apps.py | 4 +- winter/web/__init__.py | 6 +- winter/web/throttling/__init__.py | 3 +- .../web/throttling/redis_throttling_client.py | 22 +++++++ .../redis_throttling_configuration.py | 14 +++++ winter/web/throttling/throttling.py | 24 +++---- .../throttling_statistic_storage.py | 63 ------------------- 7 files changed, 52 insertions(+), 84 deletions(-) create mode 100644 winter/web/throttling/redis_throttling_client.py delete mode 100644 winter/web/throttling/throttling_statistic_storage.py diff --git a/tests/apps.py b/tests/apps.py index 434753a5..d96d5a80 100644 --- a/tests/apps.py +++ b/tests/apps.py @@ -28,6 +28,8 @@ def ready(self): winter_openapi.setup() + winter.web.setup() + self._redis_container = RedisContainer() self._redis_container.start() self._redis_container.get_client().flushdb() @@ -39,7 +41,7 @@ def ready(self): db=0, password=self._redis_container.password ) - winter.web.setup(redis_throttling_configuration) + winter.web.set_redis_throttling_configuration(redis_throttling_configuration) winter_django.setup() diff --git a/winter/web/__init__.py b/winter/web/__init__.py index 16c467d1..93d6b178 100644 --- a/winter/web/__init__.py +++ b/winter/web/__init__.py @@ -20,13 +20,13 @@ from .response_header_resolver import ResponseHeaderArgumentResolver from .response_header_serializer import response_headers_serializer from .response_status_annotation import response_status -from .throttling import init_throttling_statistic_storage from .throttling import RedisThrottlingConfiguration +from .throttling import set_redis_throttling_configuration from .throttling import throttling from .urls import register_url_regexp -def setup(redis_throttling_configuration: RedisThrottlingConfiguration): +def setup(): from winter.data.exceptions import NotFoundException from .exceptions import RedirectException from .exceptions.problem_handling import autodiscover_problem_annotations @@ -63,5 +63,3 @@ def setup(redis_throttling_configuration: RedisThrottlingConfiguration): } for exception_class, handler in auto_handle_exceptions.items(): exception_handlers_registry.add_handler(exception_class, handler, auto_handle=True) - - init_throttling_statistic_storage(redis_throttling_configuration) diff --git a/winter/web/throttling/__init__.py b/winter/web/throttling/__init__.py index fe1f47d8..63419eb1 100644 --- a/winter/web/throttling/__init__.py +++ b/winter/web/throttling/__init__.py @@ -1,6 +1,5 @@ from .throttling import throttling from .throttling import reset from .throttling import create_throttle_class -from .throttling_statistic_storage import init_throttling_statistic_storage -from .throttling_statistic_storage import get_throttling_statistic_storage +from .redis_throttling_configuration import set_redis_throttling_configuration from .redis_throttling_configuration import RedisThrottlingConfiguration diff --git a/winter/web/throttling/redis_throttling_client.py b/winter/web/throttling/redis_throttling_client.py new file mode 100644 index 00000000..14a095d7 --- /dev/null +++ b/winter/web/throttling/redis_throttling_client.py @@ -0,0 +1,22 @@ +from redis import Redis + +from .redis_throttling_configuration import get_redis_throttling_configuration + +_redis_throttling_client: Redis | None = None + +def get_redis_throttling_client() -> Redis: + global _redis_throttling_client + if _redis_throttling_client is None: + configuration = get_redis_throttling_configuration() + + if configuration is None: + raise RuntimeError(f'Configuration for Redis must be set before using the throttling') + + _redis_throttling_client = Redis( + host=configuration.host, + port=configuration.port, + db=configuration.db, + password=configuration.password, + decode_responses=True, + ) + return _redis_throttling_client diff --git a/winter/web/throttling/redis_throttling_configuration.py b/winter/web/throttling/redis_throttling_configuration.py index 5f64b41d..63ae02a1 100644 --- a/winter/web/throttling/redis_throttling_configuration.py +++ b/winter/web/throttling/redis_throttling_configuration.py @@ -7,3 +7,17 @@ class RedisThrottlingConfiguration: port: int db: int password: str | None = None + + +_redis_throttling_configuration: RedisThrottlingConfiguration | None = None + + +def set_redis_throttling_configuration(configuration: RedisThrottlingConfiguration): + global _redis_throttling_configuration + if _redis_throttling_configuration is not None: + raise RuntimeError(f'{RedisThrottlingConfiguration.__name__} is already initialized.') + _redis_throttling_configuration = configuration + + +def get_redis_throttling_configuration() -> RedisThrottlingConfiguration | None: + return _redis_throttling_configuration \ No newline at end of file diff --git a/winter/web/throttling/throttling.py b/winter/web/throttling/throttling.py index 877963ea..8e182771 100644 --- a/winter/web/throttling/throttling.py +++ b/winter/web/throttling/throttling.py @@ -7,7 +7,7 @@ import django.http from winter.core import annotate_method -from .throttling_statistic_storage import get_throttling_statistic_storage +from .redis_throttling_client import get_redis_throttling_client if TYPE_CHECKING: from winter.web.routing import Route # noqa: F401 @@ -33,25 +33,21 @@ def throttling(rate: Optional[str], scope: Optional[str] = None): class BaseRateThrottle: def __init__(self, throttling_: Throttling): self._throttling = throttling_ - self._storage = get_throttling_statistic_storage() + self._redis_client = get_redis_throttling_client() def allow_request(self, request: django.http.HttpRequest) -> bool: ident = _get_ident(request) key = _get_cache_key(self._throttling.scope, ident) + now = time.time() - with self._storage.lock(key, timeout_sec=2): - history = self._storage.get(key) - now = time.time() + pipe = self._redis_client.pipeline() + pipe.zadd(key, {str(now): now}) + pipe.zremrangebyscore(key, 0, now - self._throttling.duration) + pipe.zcard(key) + pipe.expire(key, self._throttling.duration) + _, _, count, _ = pipe.execute() - while history and history[-1] <= now - self._throttling.duration: - history.pop() - - if len(history) >= self._throttling.num_requests: - return False - - history.insert(0, now) - self._storage.set(key, history, self._throttling.duration) - return True + return count <= self._throttling.num_requests def reset(request: django.http.HttpRequest, scope: str): diff --git a/winter/web/throttling/throttling_statistic_storage.py b/winter/web/throttling/throttling_statistic_storage.py deleted file mode 100644 index b86d61be..00000000 --- a/winter/web/throttling/throttling_statistic_storage.py +++ /dev/null @@ -1,63 +0,0 @@ -from contextlib import contextmanager -from typing import Generator - -from redis import Redis -from redis.exceptions import LockError -from redis.exceptions import LockNotOwnedError - -from .redis_throttling_configuration import RedisThrottlingConfiguration - - -class ThrottlingStatisticStorage: - def __init__(self, configuration: RedisThrottlingConfiguration): - self._redis_client = Redis( - host=configuration.host, - port=configuration.port, - db=configuration.db, - password=configuration.password, - decode_responses=True, - ) - - @contextmanager - def lock(self, key: str, timeout_sec: int) -> Generator[None, None, None]: - lock = self._redis_client.lock(f'lock:{key}', timeout=timeout_sec, blocking=True) - lock.acquire() - - try: - yield - finally: - try: - lock.release() - except (LockError, LockNotOwnedError): - pass - - def get(self, key: str) -> list[float]: - raw_value: str | None = self._redis_client.get(key) - if raw_value is None: - return [] - - value = map(float, raw_value.split(',')) - return list(value) - - def set(self, key: str, value: list[float], ex: int): - raw_value = ','.join(map(str, value)) - self._redis_client.set(key, raw_value, ex) - - def delete(self, key: str): - self._redis_client.delete(key) - - -_throttling_statistic_storage: ThrottlingStatisticStorage | None = None - - -def init_throttling_statistic_storage(configuration: RedisThrottlingConfiguration): - global _throttling_statistic_storage - if _throttling_statistic_storage is not None: - raise RuntimeError("ThrottlingStatisticStorage is already initialized.") - _throttling_statistic_storage = ThrottlingStatisticStorage(configuration) - - -def get_throttling_statistic_storage() -> ThrottlingStatisticStorage: - if _throttling_statistic_storage is None: - raise RuntimeError("ThrottlingStatisticStorage is not initialized.") - return _throttling_statistic_storage \ No newline at end of file From ee08abb762f808ae34d8a146cd3739e4669ccada Mon Sep 17 00:00:00 2001 From: Andrej Romanov Date: Wed, 25 Jun 2025 22:23:43 +0300 Subject: [PATCH 3/6] rejected requests won't be counted. Avoid race conditions without using locks --- winter/web/__init__.py | 1 + winter/web/throttling/__init__.py | 1 + winter/web/throttling/exceptions.py | 2 ++ .../web/throttling/redis_throttling_client.py | 34 +++++++++++++++++-- winter/web/throttling/throttling.py | 19 +++++------ 5 files changed, 43 insertions(+), 14 deletions(-) create mode 100644 winter/web/throttling/exceptions.py diff --git a/winter/web/__init__.py b/winter/web/__init__.py index 93d6b178..e123c493 100644 --- a/winter/web/__init__.py +++ b/winter/web/__init__.py @@ -20,6 +20,7 @@ from .response_header_resolver import ResponseHeaderArgumentResolver from .response_header_serializer import response_headers_serializer from .response_status_annotation import response_status +from .throttling import ThrottlingMisconfigurationException from .throttling import RedisThrottlingConfiguration from .throttling import set_redis_throttling_configuration from .throttling import throttling diff --git a/winter/web/throttling/__init__.py b/winter/web/throttling/__init__.py index 63419eb1..6677a8d8 100644 --- a/winter/web/throttling/__init__.py +++ b/winter/web/throttling/__init__.py @@ -1,3 +1,4 @@ +from .exceptions import ThrottlingMisconfigurationException from .throttling import throttling from .throttling import reset from .throttling import create_throttle_class diff --git a/winter/web/throttling/exceptions.py b/winter/web/throttling/exceptions.py new file mode 100644 index 00000000..66e83489 --- /dev/null +++ b/winter/web/throttling/exceptions.py @@ -0,0 +1,2 @@ +class ThrottlingMisconfigurationException(Exception): + pass diff --git a/winter/web/throttling/redis_throttling_client.py b/winter/web/throttling/redis_throttling_client.py index 14a095d7..4f3e7f66 100644 --- a/winter/web/throttling/redis_throttling_client.py +++ b/winter/web/throttling/redis_throttling_client.py @@ -1,16 +1,42 @@ from redis import Redis +from redis.commands.core import Script +from .exceptions import ThrottlingMisconfigurationException from .redis_throttling_configuration import get_redis_throttling_configuration +# Redis Lua scripts are atomic +# Sliding window throttling. +# Rejected requests aren't counted. +RATE_LIMIT_LUA = ''' +local key = KEYS[1] +local now = tonumber(ARGV[1]) +local duration = tonumber(ARGV[2]) +local max_requests = tonumber(ARGV[3]) + +redis.call("ZREMRANGEBYSCORE", key, 0, now - duration) +local count = redis.call("ZCARD", key) + +if count >= max_requests then + return 0 +end + +redis.call("ZADD", key, now, now) +redis.call("EXPIRE", key, duration) +return 1 +''' + _redis_throttling_client: Redis | None = None +_rate_limit_script: Script | None = None -def get_redis_throttling_client() -> Redis: +def get_redis_throttling_client() -> tuple[Redis, Script]: global _redis_throttling_client + global _rate_limit_script + if _redis_throttling_client is None: configuration = get_redis_throttling_configuration() if configuration is None: - raise RuntimeError(f'Configuration for Redis must be set before using the throttling') + raise ThrottlingMisconfigurationException('Configuration for Redis must be set before using the throttling') _redis_throttling_client = Redis( host=configuration.host, @@ -19,4 +45,6 @@ def get_redis_throttling_client() -> Redis: password=configuration.password, decode_responses=True, ) - return _redis_throttling_client + + _rate_limit_script = _redis_throttling_client.register_script(RATE_LIMIT_LUA) + return _redis_throttling_client, _rate_limit_script diff --git a/winter/web/throttling/throttling.py b/winter/web/throttling/throttling.py index 8e182771..e9bc2848 100644 --- a/winter/web/throttling/throttling.py +++ b/winter/web/throttling/throttling.py @@ -33,21 +33,19 @@ def throttling(rate: Optional[str], scope: Optional[str] = None): class BaseRateThrottle: def __init__(self, throttling_: Throttling): self._throttling = throttling_ - self._redis_client = get_redis_throttling_client() + self._redis_client, self._rate_limit_script = get_redis_throttling_client() def allow_request(self, request: django.http.HttpRequest) -> bool: ident = _get_ident(request) key = _get_cache_key(self._throttling.scope, ident) now = time.time() - pipe = self._redis_client.pipeline() - pipe.zadd(key, {str(now): now}) - pipe.zremrangebyscore(key, 0, now - self._throttling.duration) - pipe.zcard(key) - pipe.expire(key, self._throttling.duration) - _, _, count, _ = pipe.execute() + is_allowed = self._rate_limit_script( + keys=[key], + args=[now, self._throttling.duration, self._throttling.num_requests] + ) - return count <= self._throttling.num_requests + return is_allowed == 1 def reset(request: django.http.HttpRequest, scope: str): @@ -57,9 +55,8 @@ def reset(request: django.http.HttpRequest, scope: str): """ ident = _get_ident(request) key = _get_cache_key(scope, ident) - storage = get_throttling_statistic_storage() - with storage.lock(key, timeout_sec=2): - storage.delete(key) + redis_client, _ = get_redis_throttling_client() + redis_client.delete(key) CACHE_KEY_FORMAT = 'throttle_{scope}_{ident}' From 5aaf2e6cb57f6297f25fea0a691a09df1655bbaa Mon Sep 17 00:00:00 2001 From: Andrej Romanov Date: Wed, 25 Jun 2025 22:25:32 +0300 Subject: [PATCH 4/6] get rid redundant dependency --- pyproject.toml | 1 - winter/web/throttling/redis_throttling_configuration.py | 4 +++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 95203a41..d118de3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,7 +63,6 @@ pytz = ">=2020.5" [tool.poetry.group.dev.dependencies] setuptools = "^71.1.0" testcontainers = "^4.10.0" -redis = "^6.2.0" [build-system] requires = ["poetry-core>=1.3.1"] diff --git a/winter/web/throttling/redis_throttling_configuration.py b/winter/web/throttling/redis_throttling_configuration.py index 63ae02a1..60f7353b 100644 --- a/winter/web/throttling/redis_throttling_configuration.py +++ b/winter/web/throttling/redis_throttling_configuration.py @@ -1,5 +1,7 @@ from dataclasses import dataclass +from .exceptions import ThrottlingMisconfigurationException + @dataclass class RedisThrottlingConfiguration: @@ -15,7 +17,7 @@ class RedisThrottlingConfiguration: def set_redis_throttling_configuration(configuration: RedisThrottlingConfiguration): global _redis_throttling_configuration if _redis_throttling_configuration is not None: - raise RuntimeError(f'{RedisThrottlingConfiguration.__name__} is already initialized.') + raise ThrottlingMisconfigurationException(f'{RedisThrottlingConfiguration.__name__} is already initialized.') _redis_throttling_configuration = configuration From 2e352b223363ac76385d0dde46a319ebcbefb727 Mon Sep 17 00:00:00 2001 From: Andrej Romanov Date: Wed, 25 Jun 2025 23:43:29 +0300 Subject: [PATCH 5/6] RedisThrottlingClient added --- tests/apps.py | 4 +- .../web/throttling/redis_throttling_client.py | 78 ++++++++++++------- winter/web/throttling/throttling.py | 12 +-- 3 files changed, 53 insertions(+), 41 deletions(-) diff --git a/tests/apps.py b/tests/apps.py index d96d5a80..9539f99c 100644 --- a/tests/apps.py +++ b/tests/apps.py @@ -37,7 +37,7 @@ def ready(self): redis_throttling_configuration = RedisThrottlingConfiguration( host=self._redis_container.get_container_host_ip(), - port=self._redis_container.port, + port=self._redis_container.get_exposed_port(self._redis_container.port), db=0, password=self._redis_container.password ) @@ -47,6 +47,6 @@ def ready(self): exception_handlers_registry.set_default_handler(DefaultExceptionHandler) # for 100% test coverage - def cleanup_redis(self): + def cleanup_redis(self): # pragma: no cover if self._redis_container: self._redis_container.stop() diff --git a/winter/web/throttling/redis_throttling_client.py b/winter/web/throttling/redis_throttling_client.py index 4f3e7f66..16fd6437 100644 --- a/winter/web/throttling/redis_throttling_client.py +++ b/winter/web/throttling/redis_throttling_client.py @@ -1,36 +1,61 @@ +import time + from redis import Redis from redis.commands.core import Script from .exceptions import ThrottlingMisconfigurationException from .redis_throttling_configuration import get_redis_throttling_configuration +from .redis_throttling_configuration import RedisThrottlingConfiguration + + +class RedisThrottlingClient: + # Redis Lua scripts are atomic + # Sliding window throttling. + # Rejected requests aren't counted. + THROTTLING_LUA = ''' + local key = KEYS[1] + local now = tonumber(ARGV[1]) + local duration = tonumber(ARGV[2]) + local max_requests = tonumber(ARGV[3]) -# Redis Lua scripts are atomic -# Sliding window throttling. -# Rejected requests aren't counted. -RATE_LIMIT_LUA = ''' -local key = KEYS[1] -local now = tonumber(ARGV[1]) -local duration = tonumber(ARGV[2]) -local max_requests = tonumber(ARGV[3]) + redis.call("ZREMRANGEBYSCORE", key, 0, now - duration) + local count = redis.call("ZCARD", key) -redis.call("ZREMRANGEBYSCORE", key, 0, now - duration) -local count = redis.call("ZCARD", key) + if count >= max_requests then + return 0 + end -if count >= max_requests then - return 0 -end + redis.call("ZADD", key, now, now) + redis.call("EXPIRE", key, duration) + return 1 + ''' -redis.call("ZADD", key, now, now) -redis.call("EXPIRE", key, duration) -return 1 -''' + def __init__(self, configuration: RedisThrottlingConfiguration): + self._redis_client = Redis( + host=configuration.host, + port=configuration.port, + db=configuration.db, + password=configuration.password, + decode_responses=True, + ) + self._throttling_script = self._redis_client.register_script(self.THROTTLING_LUA) -_redis_throttling_client: Redis | None = None -_rate_limit_script: Script | None = None + def is_request_allowed(self, key: str, duration: int, num_requests: int) -> bool: + now = time.time() + is_allowed = self._throttling_script( + keys=[key], + args=[now, duration, num_requests] + ) + return is_allowed == 1 -def get_redis_throttling_client() -> tuple[Redis, Script]: + def delete(self, key: str): + self._redis_client.delete(key) + + +_redis_throttling_client: RedisThrottlingClient | None = None + +def get_redis_throttling_client() -> RedisThrottlingClient: global _redis_throttling_client - global _rate_limit_script if _redis_throttling_client is None: configuration = get_redis_throttling_configuration() @@ -38,13 +63,6 @@ def get_redis_throttling_client() -> tuple[Redis, Script]: if configuration is None: raise ThrottlingMisconfigurationException('Configuration for Redis must be set before using the throttling') - _redis_throttling_client = Redis( - host=configuration.host, - port=configuration.port, - db=configuration.db, - password=configuration.password, - decode_responses=True, - ) + _redis_throttling_client = RedisThrottlingClient(configuration) - _rate_limit_script = _redis_throttling_client.register_script(RATE_LIMIT_LUA) - return _redis_throttling_client, _rate_limit_script + return _redis_throttling_client diff --git a/winter/web/throttling/throttling.py b/winter/web/throttling/throttling.py index e9bc2848..0b0fd6d7 100644 --- a/winter/web/throttling/throttling.py +++ b/winter/web/throttling/throttling.py @@ -33,19 +33,13 @@ def throttling(rate: Optional[str], scope: Optional[str] = None): class BaseRateThrottle: def __init__(self, throttling_: Throttling): self._throttling = throttling_ - self._redis_client, self._rate_limit_script = get_redis_throttling_client() + self._redis_client = get_redis_throttling_client() def allow_request(self, request: django.http.HttpRequest) -> bool: ident = _get_ident(request) key = _get_cache_key(self._throttling.scope, ident) - now = time.time() - is_allowed = self._rate_limit_script( - keys=[key], - args=[now, self._throttling.duration, self._throttling.num_requests] - ) - - return is_allowed == 1 + return self._redis_client.is_request_allowed(key, self._throttling.duration, self._throttling.num_requests) def reset(request: django.http.HttpRequest, scope: str): @@ -55,7 +49,7 @@ def reset(request: django.http.HttpRequest, scope: str): """ ident = _get_ident(request) key = _get_cache_key(scope, ident) - redis_client, _ = get_redis_throttling_client() + redis_client = get_redis_throttling_client() redis_client.delete(key) From 0968f56b950f74e5b36b69775f19338fee6221d1 Mon Sep 17 00:00:00 2001 From: Andrej Romanov Date: Thu, 26 Jun 2025 17:07:31 +0300 Subject: [PATCH 6/6] additional tests added --- tests/web/test_throttling.py | 39 +++++++++++++++++++ .../web/throttling/redis_throttling_client.py | 1 - .../redis_throttling_configuration.py | 4 +- 3 files changed, 41 insertions(+), 3 deletions(-) diff --git a/tests/web/test_throttling.py b/tests/web/test_throttling.py index c448c1a9..ca3c5cb8 100644 --- a/tests/web/test_throttling.py +++ b/tests/web/test_throttling.py @@ -3,7 +3,14 @@ import freezegun import pytest +from mock import patch +from winter.web import RedisThrottlingConfiguration +from winter.web import ThrottlingMisconfigurationException +from winter.web import set_redis_throttling_configuration +from winter.web.throttling.redis_throttling_client import get_redis_throttling_client +from winter.web.throttling import redis_throttling_client +from winter.web.throttling import redis_throttling_configuration expected_error_response = { 'status': 429, @@ -65,3 +72,35 @@ def test_get_throttling_with_conditional_reset(api_client): is_reset = True if i == 5 else False response = api_client.get(f'/with-throttling/with-reset/?is_reset={is_reset}') assert response.status_code == HTTPStatus.OK, i + + +@patch.object(redis_throttling_client, 'get_redis_throttling_configuration', return_value=None) +@patch.object(redis_throttling_client, '_redis_throttling_client', None) +def test_get_redis_throttling_client_without_configuration(_): + with pytest.raises(ThrottlingMisconfigurationException) as exc_info: + get_redis_throttling_client() + + assert 'Configuration for Redis must be set' in str(exc_info.value) + + +@patch.object( + redis_throttling_configuration, + '_redis_throttling_configuration', + RedisThrottlingConfiguration( + host='localhost', + port=1234, + db=0, + password=None + ) +) +def test_try_to_set_redis_configuration_twice(): + configuration = RedisThrottlingConfiguration( + host='localhost', + port=5678, + db=0, + password=None + ) + with pytest.raises(ThrottlingMisconfigurationException) as exc_info: + set_redis_throttling_configuration(configuration) + + assert 'RedisThrottlingConfiguration is already initialized' in str(exc_info.value) diff --git a/winter/web/throttling/redis_throttling_client.py b/winter/web/throttling/redis_throttling_client.py index 16fd6437..a99122ea 100644 --- a/winter/web/throttling/redis_throttling_client.py +++ b/winter/web/throttling/redis_throttling_client.py @@ -1,7 +1,6 @@ import time from redis import Redis -from redis.commands.core import Script from .exceptions import ThrottlingMisconfigurationException from .redis_throttling_configuration import get_redis_throttling_configuration diff --git a/winter/web/throttling/redis_throttling_configuration.py b/winter/web/throttling/redis_throttling_configuration.py index 60f7353b..410c3d63 100644 --- a/winter/web/throttling/redis_throttling_configuration.py +++ b/winter/web/throttling/redis_throttling_configuration.py @@ -17,9 +17,9 @@ class RedisThrottlingConfiguration: def set_redis_throttling_configuration(configuration: RedisThrottlingConfiguration): global _redis_throttling_configuration if _redis_throttling_configuration is not None: - raise ThrottlingMisconfigurationException(f'{RedisThrottlingConfiguration.__name__} is already initialized.') + raise ThrottlingMisconfigurationException(f'{RedisThrottlingConfiguration.__name__} is already initialized') _redis_throttling_configuration = configuration def get_redis_throttling_configuration() -> RedisThrottlingConfiguration | None: - return _redis_throttling_configuration \ No newline at end of file + return _redis_throttling_configuration