limits 4.0.1__py3-none-any.whl → 4.1__py3-none-any.whl

limits/aio/storage/memory.py

@@ -1,12 +1,18 @@
 import asyncio
 import time
-from collections import Counter
+from collections import Counter, defaultdict
+from math import floor
 
 from deprecated.sphinx import versionadded
 
 import limits.typing
-from limits.aio.storage.base import MovingWindowSupport, Storage
-from limits.typing import Dict, List, Optional, Tuple, Type, Union
+from limits.aio.storage.base import (
+    MovingWindowSupport,
+    SlidingWindowCounterSupport,
+    Storage,
+)
+from limits.storage.base import TimestampedSlidingWindow
+from limits.typing import Optional, Type, Union
 
 
 class LockableEntry(asyncio.Lock):
@@ -17,7 +23,9 @@ class LockableEntry(asyncio.Lock):
 
 
 @versionadded(version="2.1")
-class MemoryStorage(Storage, MovingWindowSupport):
+class MemoryStorage(
+    Storage, MovingWindowSupport, SlidingWindowCounterSupport, TimestampedSlidingWindow
+):
     """
     rate limit storage using :class:`collections.Counter`
     as an in memory storage for fixed and elastic window strategies,
@@ -34,15 +42,16 @@ class MemoryStorage(Storage, MovingWindowSupport):
         self, uri: Optional[str] = None, wrap_exceptions: bool = False, **_: str
     ) -> None:
         self.storage: limits.typing.Counter[str] = Counter()
-        self.expirations: Dict[str, float] = {}
-        self.events: Dict[str, List[LockableEntry]] = {}
+        self.locks: defaultdict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
+        self.expirations: dict[str, float] = {}
+        self.events: dict[str, list[LockableEntry]] = {}
         self.timer: Optional[asyncio.Task[None]] = None
         super().__init__(uri, wrap_exceptions=wrap_exceptions, **_)
 
     @property
     def base_exceptions(
         self,
-    ) -> Union[Type[Exception], Tuple[Type[Exception], ...]]:  # pragma: no cover
+    ) -> Union[Type[Exception], tuple[Type[Exception], ...]]:  # pragma: no cover
         return ValueError
 
     async def __expire_events(self) -> None:
@@ -56,13 +65,14 @@ class MemoryStorage(Storage, MovingWindowSupport):
             if self.expirations[key] <= time.time():
                 self.storage.pop(key, None)
                 self.expirations.pop(key, None)
+                self.locks.pop(key, None)
 
     async def __schedule_expiry(self) -> None:
         if not self.timer or self.timer.done():
             self.timer = asyncio.create_task(self.__expire_events())
 
     async def incr(
-        self, key: str, expiry: int, elastic_expiry: bool = False, amount: int = 1
+        self, key: str, expiry: float, elastic_expiry: bool = False, amount: int = 1
     ) -> int:
         """
         increments the counter for a given rate limit key
@@ -75,10 +85,24 @@ class MemoryStorage(Storage, MovingWindowSupport):
         """
         await self.get(key)
         await self.__schedule_expiry()
-        self.storage[key] += amount
+        async with self.locks[key]:
+            self.storage[key] += amount
 
-        if elastic_expiry or self.storage[key] == amount:
-            self.expirations[key] = time.time() + expiry
+            if elastic_expiry or self.storage[key] == amount:
+                self.expirations[key] = time.time() + expiry
+
+        return self.storage.get(key, amount)
+
+    async def decr(self, key: str, amount: int = 1) -> int:
+        """
+        decrements the counter for a given rate limit key. 0 is the minimum allowed value.
+
+        :param amount: the number to increment by
+        """
+        await self.get(key)
+        await self.__schedule_expiry()
+        async with self.locks[key]:
+            self.storage[key] = max(self.storage[key] - amount, 0)
 
         return self.storage.get(key, amount)
 
@@ -86,10 +110,10 @@ class MemoryStorage(Storage, MovingWindowSupport):
         """
         :param key: the key to get the counter value for
         """
-
         if self.expirations.get(key, 0) <= time.time():
             self.storage.pop(key, None)
             self.expirations.pop(key, None)
+            self.locks.pop(key, None)
 
         return self.storage.get(key, 0)
 
@@ -100,6 +124,7 @@ class MemoryStorage(Storage, MovingWindowSupport):
         self.storage.pop(key, None)
         self.expirations.pop(key, None)
         self.events.pop(key, None)
+        self.locks.pop(key, None)
 
     async def acquire_entry(
         self, key: str, limit: int, expiry: int, amount: int = 1
@@ -153,7 +178,7 @@ class MemoryStorage(Storage, MovingWindowSupport):
     # FIXME: arg limit is not used
     async def get_moving_window(
         self, key: str, limit: int, expiry: int
-    ) -> Tuple[float, int]:
+    ) -> tuple[float, int]:
         """
         returns the starting point and the number of entries in the moving
         window
@@ -171,6 +196,65 @@ class MemoryStorage(Storage, MovingWindowSupport):
 
         return timestamp, acquired
 
+    async def acquire_sliding_window_entry(
+        self,
+        key: str,
+        limit: int,
+        expiry: int,
+        amount: int = 1,
+    ) -> bool:
+        if amount > limit:
+            return False
+        now = time.time()
+        previous_key, current_key = self.sliding_window_keys(key, expiry, now)
+        (
+            previous_count,
+            previous_ttl,
+            current_count,
+            _,
+        ) = await self._get_sliding_window_info(previous_key, current_key, expiry, now)
+        weighted_count = previous_count * previous_ttl / expiry + current_count
+        if floor(weighted_count) + amount > limit:
+            return False
+        else:
+            # Hit, increase the current counter.
+            # If the counter doesn't exist yet, set twice the theorical expiry.
+            current_count = await self.incr(current_key, 2 * expiry, amount=amount)
+            weighted_count = previous_count * previous_ttl / expiry + current_count
+            if floor(weighted_count) > limit:
+                # Another hit won the race condition: revert the incrementation and refuse this hit
+                # Limitation: during high concurrency at the end of the window,
+                # the counter is shifted and cannot be decremented, so less requests than expected are allowed.
+                await self.decr(current_key, amount)
+                # print("Concurrent call, reverting the counter increase")
+                return False
+            return True
+
+    async def get_sliding_window(
+        self, key: str, expiry: int
+    ) -> tuple[int, float, int, float]:
+        now = time.time()
+        previous_key, current_key = self.sliding_window_keys(key, expiry, now)
+        return await self._get_sliding_window_info(
+            previous_key, current_key, expiry, now
+        )
+
+    async def _get_sliding_window_info(
+        self,
+        previous_key: str,
+        current_key: str,
+        expiry: int,
+        now: float,
+    ) -> tuple[int, float, int, float]:
+        previous_count = await self.get(previous_key)
+        current_count = await self.get(current_key)
+        if previous_count == 0:
+            previous_ttl = float(0)
+        else:
+            previous_ttl = (1 - (((now - expiry) / expiry) % 1)) * expiry
+        current_ttl = (1 - ((now / expiry) % 1)) * expiry + expiry
+        return previous_count, previous_ttl, current_count, current_ttl
+
     async def check(self) -> bool:
         """
         check if storage is healthy
@@ -183,5 +267,6 @@ class MemoryStorage(Storage, MovingWindowSupport):
         self.storage.clear()
         self.expirations.clear()
         self.events.clear()
+        self.locks.clear()
 
         return num_items

limits/aio/storage/mongodb.py

@@ -3,12 +3,22 @@ from __future__ import annotations
 import asyncio
 import datetime
 import time
-from typing import cast
 
 from deprecated.sphinx import versionadded, versionchanged
 
-from limits.aio.storage.base import MovingWindowSupport, Storage
-from limits.typing import Dict, List, Optional, ParamSpec, Tuple, Type, TypeVar, Union
+from limits.aio.storage.base import (
+    MovingWindowSupport,
+    SlidingWindowCounterSupport,
+    Storage,
+)
+from limits.typing import (
+    Optional,
+    ParamSpec,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
 from limits.util import get_dependency
 
 P = ParamSpec("P")
@@ -20,7 +30,7 @@ R = TypeVar("R")
     version="3.14.0",
     reason="Added option to select custom collection names for windows & counters",
 )
-class MongoDBStorage(Storage, MovingWindowSupport):
+class MongoDBStorage(Storage, MovingWindowSupport, SlidingWindowCounterSupport):
     """
     Rate limit storage with MongoDB as backend.
 
@@ -50,7 +60,8 @@ class MongoDBStorage(Storage, MovingWindowSupport):
          collections.
         :param counter_collection_name: The collection name to use for individual counters
          used in fixed window strategies
-        :param window_collection_name: The collection name to use for moving window storage
+        :param window_collection_name: The collection name to use for sliding & moving window
+         storage
         :param wrap_exceptions: Whether to wrap storage exceptions in
          :exc:`limits.errors.StorageError` before raising it.
         :param options: all remaining keyword arguments are passed
@@ -83,7 +94,7 @@ class MongoDBStorage(Storage, MovingWindowSupport):
     @property
     def base_exceptions(
         self,
-    ) -> Union[Type[Exception], Tuple[Type[Exception], ...]]:  # pragma: no cover
+    ) -> Union[Type[Exception], tuple[Type[Exception], ...]]:  # pragma: no cover
         return self.lib_errors.PyMongoError  # type: ignore
 
     @property
@@ -224,7 +235,7 @@ class MongoDBStorage(Storage, MovingWindowSupport):
 
     async def get_moving_window(
         self, key: str, limit: int, expiry: int
-    ) -> Tuple[float, int]:
+    ) -> tuple[float, int]:
         """
         returns the starting point and the number of entries in the moving
         window
@@ -281,9 +292,9 @@ class MongoDBStorage(Storage, MovingWindowSupport):
 
         timestamp = time.time()
         try:
-            updates: Dict[
+            updates: dict[
                 str,
-                Dict[str, Union[datetime.datetime, Dict[str, Union[List[float], int]]]],
+                dict[str, Union[datetime.datetime, dict[str, Union[list[float], int]]]],
             ] = {
                 "$push": {
                     "entries": {
@@ -312,3 +323,200 @@ class MongoDBStorage(Storage, MovingWindowSupport):
             return True
         except self.proxy_dependency.module.errors.DuplicateKeyError:
             return False
+
+    async def acquire_sliding_window_entry(
+        self, key: str, limit: int, expiry: int, amount: int = 1
+    ) -> bool:
+        await self.create_indices()
+        expiry_ms = expiry * 1000
+        result = await self.database[
+            self.__collection_mapping["windows"]
+        ].find_one_and_update(
+            {"_id": key},
+            [
+                {
+                    "$set": {
+                        "previousCount": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": {"$ifNull": ["$currentCount", 0]},
+                                "else": {"$ifNull": ["$previousCount", 0]},
+                            }
+                        },
+                    }
+                },
+                {
+                    "$set": {
+                        "currentCount": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": 0,
+                                "else": {"$ifNull": ["$currentCount", 0]},
+                            }
+                        },
+                        "expiresAt": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": {
+                                    "$cond": {
+                                        "if": {"$gt": ["$expiresAt", 0]},
+                                        "then": {"$add": ["$expiresAt", expiry_ms]},
+                                        "else": {"$add": ["$$NOW", 2 * expiry_ms]},
+                                    }
+                                },
+                                "else": "$expiresAt",
+                            }
+                        },
+                    }
+                },
+                {
+                    "$set": {
+                        "curWeightedCount": {
+                            "$floor": {
+                                "$add": [
+                                    {
+                                        "$multiply": [
+                                            "$previousCount",
+                                            {
+                                                "$divide": [
+                                                    {
+                                                        "$max": [
+                                                            0,
+                                                            {
+                                                                "$subtract": [
+                                                                    "$expiresAt",
+                                                                    {
+                                                                        "$add": [
+                                                                            "$$NOW",
+                                                                            expiry_ms,
+                                                                        ]
+                                                                    },
+                                                                ]
+                                                            },
+                                                        ]
+                                                    },
+                                                    expiry_ms,
+                                                ]
+                                            },
+                                        ]
+                                    },
+                                    "$currentCount",
+                                ]
+                            }
+                        }
+                    }
+                },
+                {
+                    "$set": {
+                        "currentCount": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$add": ["$curWeightedCount", amount]},
+                                        limit,
+                                    ]
+                                },
+                                "then": {"$add": ["$currentCount", amount]},
+                                "else": "$currentCount",
+                            }
+                        }
+                    }
+                },
+                {
+                    "$set": {
+                        "_acquired": {
+                            "$lte": [{"$add": ["$curWeightedCount", amount]}, limit]
+                        }
+                    }
+                },
+                {"$unset": ["curWeightedCount"]},
+            ],
+            return_document=self.proxy_dependency.module.ReturnDocument.AFTER,
+            upsert=True,
+        )
+
+        return cast(bool, result["_acquired"])
+
+    async def get_sliding_window(
+        self, key: str, expiry: int
+    ) -> tuple[int, float, int, float]:
+        expiry_ms = expiry * 1000
+        if result := await self.database[
+            self.__collection_mapping["windows"]
+        ].find_one_and_update(
+            {"_id": key},
+            [
+                {
+                    "$set": {
+                        "previousCount": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": {"$ifNull": ["$currentCount", 0]},
+                                "else": {"$ifNull": ["$previousCount", 0]},
+                            }
+                        },
+                        "currentCount": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": 0,
+                                "else": {"$ifNull": ["$currentCount", 0]},
+                            }
+                        },
+                        "expiresAt": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": {"$add": ["$expiresAt", expiry_ms]},
+                                "else": "$expiresAt",
+                            }
+                        },
+                    }
+                }
+            ],
+            return_document=self.proxy_dependency.module.ReturnDocument.AFTER,
+            projection=["currentCount", "previousCount", "expiresAt"],
+        ):
+            expires_at = (
+                (result["expiresAt"].replace(tzinfo=datetime.timezone.utc).timestamp())
+                if result.get("expiresAt")
+                else time.time()
+            )
+            current_ttl = max(0, expires_at - time.time())
+            prev_ttl = max(0, current_ttl - expiry if result["previousCount"] else 0)
+
+            return (
+                result["previousCount"],
+                prev_ttl,
+                result["currentCount"],
+                current_ttl,
+            )
+        return 0, 0.0, 0, 0.0

limits/aio/storage/redis.py

@@ -5,9 +5,13 @@ from typing import TYPE_CHECKING, cast
 from deprecated.sphinx import versionadded
 from packaging.version import Version
 
-from limits.aio.storage.base import MovingWindowSupport, Storage
+from limits.aio.storage.base import (
+    MovingWindowSupport,
+    SlidingWindowCounterSupport,
+    Storage,
+)
 from limits.errors import ConfigurationError
-from limits.typing import AsyncRedisClient, Dict, Optional, Tuple, Type, Union
+from limits.typing import AsyncRedisClient, Optional, Type, Union
 from limits.util import get_package_data
 
 if TYPE_CHECKING:
@@ -24,9 +28,15 @@ class RedisInteractor:
     )
     SCRIPT_CLEAR_KEYS = get_package_data(f"{RES_DIR}/clear_keys.lua")
     SCRIPT_INCR_EXPIRE = get_package_data(f"{RES_DIR}/incr_expire.lua")
+    SCRIPT_SLIDING_WINDOW = get_package_data(f"{RES_DIR}/sliding_window.lua")
+    SCRIPT_ACQUIRE_SLIDING_WINDOW = get_package_data(
+        f"{RES_DIR}/acquire_sliding_window.lua"
+    )
 
     lua_moving_window: "coredis.commands.Script[bytes]"
-    lua_acquire_window: "coredis.commands.Script[bytes]"
+    lua_acquire_moving_window: "coredis.commands.Script[bytes]"
+    lua_sliding_window: "coredis.commands.Script[bytes]"
+    lua_acquire_sliding_window: "coredis.commands.Script[bytes]"
     lua_clear_keys: "coredis.commands.Script[bytes]"
     lua_incr_expire: "coredis.commands.Script[bytes]"
 
@@ -78,14 +88,14 @@ class RedisInteractor:
 
     async def get_moving_window(
         self, key: str, limit: int, expiry: int
-    ) -> Tuple[float, int]:
+    ) -> tuple[float, int]:
         """
         returns the starting point and the number of entries in the moving
         window
 
         :param key: rate limit key
         :param expiry: expiry of entry
-        :return: (start of window, number of acquired entries)
+        :return: (previous count, previous TTL, current count, current TTL)
         """
         key = self.prefixed_key(key)
         timestamp = time.time()
@@ -96,6 +106,23 @@ class RedisInteractor:
             return float(window[0]), window[1]  # type: ignore
         return timestamp, 0
 
+    async def get_sliding_window(
+        self, key: str, expiry: int
+    ) -> tuple[int, float, int, float]:
+        previous_key = self.prefixed_key(self._previous_window_key(key))
+        current_key = self.prefixed_key(self._current_window_key(key))
+
+        if window := await self.lua_sliding_window.execute(
+            [previous_key, current_key], [expiry]
+        ):
+            return (
+                int(window[0] or 0),  # type: ignore
+                max(0, float(window[1] or 0)) / 1000,  # type: ignore
+                int(window[2] or 0),  # type: ignore
+                max(0, float(window[3] or 0)) / 1000,  # type: ignore
+            )
+        return 0, 0.0, 0, 0.0
+
     async def _acquire_entry(
         self,
         key: str,
@@ -112,12 +139,27 @@ class RedisInteractor:
         """
         key = self.prefixed_key(key)
         timestamp = time.time()
-        acquired = await self.lua_acquire_window.execute(
+        acquired = await self.lua_acquire_moving_window.execute(
             [key], [timestamp, limit, expiry, amount]
         )
 
         return bool(acquired)
 
+    async def _acquire_sliding_window_entry(
+        self,
+        previous_key: str,
+        current_key: str,
+        limit: int,
+        expiry: int,
+        amount: int = 1,
+    ) -> bool:
+        previous_key = self.prefixed_key(previous_key)
+        current_key = self.prefixed_key(current_key)
+        acquired = await self.lua_acquire_sliding_window.execute(
+            [previous_key, current_key], [limit, expiry, amount]
+        )
+        return bool(acquired)
+
     async def _get_expiry(self, key: str, connection: AsyncRedisClient) -> float:
         """
         :param key: the key to get the expiry for
@@ -140,9 +182,35 @@ class RedisInteractor:
         except:  # noqa
             return False
 
+    def _current_window_key(self, key: str) -> str:
+        """
+        Return the current window's storage key (Sliding window strategy)
+
+        Contrary to other strategies that have one key per rate limit item,
+        this strategy has two keys per rate limit item than must be on the same machine.
+        To keep the current key and the previous key on the same Redis cluster node,
+        curly braces are added.
+
+        Eg: "{constructed_key}"
+        """
+        return f"{{{key}}}"
+
+    def _previous_window_key(self, key: str) -> str:
+        """
+        Return the previous window's storage key (Sliding window strategy).
+
+        Curvy braces are added on the common pattern with the current window's key,
+        so the current and the previous key are stored on the same Redis cluster node.
+
+        Eg: "{constructed_key}/-1"
+        """
+        return f"{self._current_window_key(key)}/-1"
+
 
 @versionadded(version="2.1")
-class RedisStorage(RedisInteractor, Storage, MovingWindowSupport):
+class RedisStorage(
+    RedisInteractor, Storage, MovingWindowSupport, SlidingWindowCounterSupport
+):
     """
     Rate limit storage with redis as backend.
 
@@ -200,18 +268,22 @@ class RedisStorage(RedisInteractor, Storage, MovingWindowSupport):
     @property
     def base_exceptions(
         self,
-    ) -> Union[Type[Exception], Tuple[Type[Exception], ...]]:  # pragma: no cover
+    ) -> Union[Type[Exception], tuple[Type[Exception], ...]]:  # pragma: no cover
         return self.dependency.exceptions.RedisError  # type: ignore[no-any-return]
 
     def initialize_storage(self, _uri: str) -> None:
         # all these methods are coroutines, so must be called with await
         self.lua_moving_window = self.storage.register_script(self.SCRIPT_MOVING_WINDOW)
-        self.lua_acquire_window = self.storage.register_script(
+        self.lua_acquire_moving_window = self.storage.register_script(
             self.SCRIPT_ACQUIRE_MOVING_WINDOW
         )
         self.lua_clear_keys = self.storage.register_script(self.SCRIPT_CLEAR_KEYS)
-        self.lua_incr_expire = self.storage.register_script(
-            RedisStorage.SCRIPT_INCR_EXPIRE
+        self.lua_incr_expire = self.storage.register_script(self.SCRIPT_INCR_EXPIRE)
+        self.lua_sliding_window = self.storage.register_script(
+            self.SCRIPT_SLIDING_WINDOW
+        )
+        self.lua_acquire_sliding_window = self.storage.register_script(
+            self.SCRIPT_ACQUIRE_SLIDING_WINDOW
         )
 
     async def incr(
@@ -261,6 +333,19 @@ class RedisStorage(RedisInteractor, Storage, MovingWindowSupport):
 
         return await super()._acquire_entry(key, limit, expiry, self.storage, amount)
 
+    async def acquire_sliding_window_entry(
+        self,
+        key: str,
+        limit: int,
+        expiry: int,
+        amount: int = 1,
+    ) -> bool:
+        current_key = self._current_window_key(key)
+        previous_key = self._previous_window_key(key)
+        return await super()._acquire_sliding_window_entry(
+            previous_key, current_key, limit, expiry, amount
+        )
+
     async def get_expiry(self, key: str) -> float:
         """
         :param key: the key to get the expiry for
@@ -302,7 +387,7 @@ class RedisClusterStorage(RedisStorage):
     The storage schemes for redis cluster to be used in an async context
     """
 
-    DEFAULT_OPTIONS: Dict[str, Union[float, str, bool]] = {
+    DEFAULT_OPTIONS: dict[str, Union[float, str, bool]] = {
         "max_connections": 1000,
     }
     "Default options passed to :class:`coredis.RedisCluster`"
@@ -322,7 +407,7 @@ class RedisClusterStorage(RedisStorage):
          available or if the redis host cannot be pinged.
         """
         parsed = urllib.parse.urlparse(uri)
-        parsed_auth: Dict[str, Union[float, str, bool]] = {}
+        parsed_auth: dict[str, Union[float, str, bool]] = {}
 
         if parsed.username:
             parsed_auth["username"] = parsed.username
@@ -386,7 +471,7 @@ class RedisSentinelStorage(RedisStorage):
         uri: str,
         service_name: Optional[str] = None,
         use_replicas: bool = True,
-        sentinel_kwargs: Optional[Dict[str, Union[float, str, bool]]] = None,
+        sentinel_kwargs: Optional[dict[str, Union[float, str, bool]]] = None,
         **options: Union[float, str, bool],
     ):
         """
@@ -407,7 +492,7 @@ class RedisSentinelStorage(RedisStorage):
         sentinel_configuration = []
         connection_options = options.copy()
         sentinel_options = sentinel_kwargs.copy() if sentinel_kwargs else {}
-        parsed_auth: Dict[str, Union[float, str, bool]] = {}
+        parsed_auth: dict[str, Union[float, str, bool]] = {}
 
         if parsed.username:
             parsed_auth["username"] = parsed.username