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

limits/storage/base.py

@@ -3,30 +3,32 @@ from __future__ import annotations
 import functools
 import threading
 from abc import ABC, abstractmethod
-from typing import Any, cast
 
 from limits import errors
 from limits.storage.registry import StorageRegistry
 from limits.typing import (
+    Any,
     Callable,
-    List,
     Optional,
     P,
     R,
-    Tuple,
     Type,
     Union,
+    cast,
 )
 from limits.util import LazyDependency
 
 
-def _wrap_errors(storage: Storage, fn: Callable[P, R]) -> Callable[P, R]:
+def _wrap_errors(
+    fn: Callable[P, R],
+) -> Callable[P, R]:
     @functools.wraps(fn)
     def inner(*args: P.args, **kwargs: P.kwargs) -> R:
+        instance = cast(Storage, args[0])
         try:
             return fn(*args, **kwargs)
-        except storage.base_exceptions as exc:
-            if storage.wrap_exceptions:
+        except instance.base_exceptions as exc:
+            if instance.wrap_exceptions:
                 raise errors.StorageError(exc) from exc
             raise
 
@@ -38,12 +40,10 @@ class Storage(LazyDependency, metaclass=StorageRegistry):
     Base class to extend when implementing a storage backend.
     """
 
-    STORAGE_SCHEME: Optional[List[str]]
+    STORAGE_SCHEME: Optional[list[str]]
     """The storage schemes to register against this implementation"""
 
-    def __new__(cls, *args: Any, **kwargs: Any) -> Storage:  # type: ignore[misc]
-        inst = super().__new__(cls)
-
+    def __init_subclass__(cls, **kwargs: Any) -> None:  # type: ignore[explicit-any]
         for method in {
             "incr",
             "get",
@@ -52,9 +52,8 @@ class Storage(LazyDependency, metaclass=StorageRegistry):
             "reset",
             "clear",
         }:
-            setattr(inst, method, _wrap_errors(inst, getattr(inst, method)))
-
-        return inst
+            setattr(cls, method, _wrap_errors(getattr(cls, method)))
+        super().__init_subclass__(**kwargs)
 
     def __init__(
         self,
@@ -73,7 +72,7 @@ class Storage(LazyDependency, metaclass=StorageRegistry):
 
     @property
     @abstractmethod
-    def base_exceptions(self) -> Union[Type[Exception], Tuple[Type[Exception], ...]]:
+    def base_exceptions(self) -> Union[Type[Exception], tuple[Type[Exception], ...]]:
         raise NotImplementedError
 
     @abstractmethod
@@ -131,24 +130,21 @@ class Storage(LazyDependency, metaclass=StorageRegistry):
 
 class MovingWindowSupport(ABC):
     """
-    Abstract base for storages that intend to support
-    the moving window strategy
+    Abstract base class for storages that support
+    the :ref:`strategies:moving window` strategy
     """
 
-    def __new__(cls, *args: Any, **kwargs: Any) -> MovingWindowSupport:  # type: ignore[misc]
-        inst = super().__new__(cls)
-
+    def __init_subclass__(cls, **kwargs: Any) -> None:  # type: ignore[explicit-any]
         for method in {
             "acquire_entry",
             "get_moving_window",
         }:
             setattr(
-                inst,
+                cls,
                 method,
-                _wrap_errors(cast(Storage, inst), getattr(inst, method)),
+                _wrap_errors(getattr(cls, method)),
             )
-
-        return inst
+        super().__init_subclass__(**kwargs)
 
     @abstractmethod
     def acquire_entry(self, key: str, limit: int, expiry: int, amount: int = 1) -> bool:
@@ -161,7 +157,7 @@ class MovingWindowSupport(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    def get_moving_window(self, key: str, limit: int, expiry: int) -> Tuple[float, int]:
+    def get_moving_window(self, key: str, limit: int, expiry: int) -> tuple[float, int]:
         """
         returns the starting point and the number of entries in the moving
         window
@@ -171,3 +167,75 @@ class MovingWindowSupport(ABC):
         :return: (start of window, number of acquired entries)
         """
         raise NotImplementedError
+
+
+class SlidingWindowCounterSupport(ABC):
+    """
+    Abstract base class for storages that support
+    the :ref:`strategies:sliding window counter` strategy.
+    """
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:  # type: ignore[explicit-any]
+        for method in {"acquire_sliding_window_entry", "get_sliding_window"}:
+            setattr(
+                cls,
+                method,
+                _wrap_errors(getattr(cls, method)),
+            )
+        super().__init_subclass__(**kwargs)
+
+    @abstractmethod
+    def acquire_sliding_window_entry(
+        self, key: str, limit: int, expiry: int, amount: int = 1
+    ) -> bool:
+        """
+        Acquire an entry if the weighted count of the current and previous
+        windows is less than or equal to the limit
+
+        :param key: rate limit key to acquire an entry in
+        :param limit: amount of entries allowed
+        :param expiry: expiry of the entry
+        :param amount: the number of entries to acquire
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_sliding_window(
+        self, key: str, expiry: int
+    ) -> tuple[int, float, int, float]:
+        """
+        Return the previous and current window information.
+
+        :param key: the rate limit key
+        :param expiry: the rate limit expiry, needed to compute the key in some implementations
+        :return: a tuple of (int, float, int, float) with the following information:
+          - previous window counter
+          - previous window TTL
+          - current window counter
+          - current window TTL
+        """
+        raise NotImplementedError
+
+
+class TimestampedSlidingWindow:
+    """Helper class for storage that support the sliding window counter, with timestamp based keys."""
+
+    @classmethod
+    def sliding_window_keys(cls, key: str, expiry: int, at: float) -> tuple[str, str]:
+        """
+        returns the previous and the current window's keys.
+
+        :param key: the key to get the window's keys from
+        :param expiry: the expiry of the limit item, in seconds
+        :param at: the timestamp to get the keys from. Default to now, ie ``time.time()``
+
+        Returns a tuple with the previous and the current key: (previous, current).
+
+        Example:
+          - key = "mykey"
+          - expiry = 60
+          - at = 1738576292.6631825
+
+        The return value will be the tuple ``("mykey/28976271", "mykey/28976270")``.
+        """
+        return f"{key}/{int((at - expiry) / expiry)}", f"{key}/{int(at / expiry)}"

limits/storage/etcd.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import time
 import urllib.parse
-from typing import TYPE_CHECKING, Optional, Tuple, Type, Union
 
 from limits.errors import ConcurrentUpdateError
 from limits.storage.base import Storage
+from limits.typing import TYPE_CHECKING, Optional, Union
 
 if TYPE_CHECKING:
     import etcd3
@@ -26,6 +28,7 @@ class EtcdStorage(Storage):
         self,
         uri: str,
         max_retries: int = MAX_RETRIES,
+        wrap_exceptions: bool = False,
         **options: str,
     ) -> None:
         """
@@ -33,6 +36,8 @@ class EtcdStorage(Storage):
          ``etcd://host:port``,
         :param max_retries: Maximum number of attempts to retry
          in the case of concurrent updates to a rate limit key
+        :param wrap_exceptions: Whether to wrap storage exceptions in
+         :exc:`limits.errors.StorageError` before raising it.
         :param options: all remaining keyword arguments are passed
          directly to the constructor of :class:`etcd3.Etcd3Client`
         :raise ConfigurationError: when :pypi:`etcd3` is not available
@@ -43,11 +48,12 @@ class EtcdStorage(Storage):
             parsed.hostname, parsed.port, **options
         )
         self.max_retries = max_retries
+        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 self.lib.Etcd3Exception  # type: ignore[no-any-return]
 
     def prefixed_key(self, key: str) -> bytes:

limits/storage/memcached.py

@@ -1,27 +1,34 @@
+from __future__ import annotations
+
 import inspect
 import threading
 import time
 import urllib.parse
+from collections.abc import Iterable
+from math import ceil, floor
 from types import ModuleType
-from typing import cast
 
 from limits.errors import ConfigurationError
-from limits.storage.base import Storage
+from limits.storage.base import (
+    SlidingWindowCounterSupport,
+    Storage,
+    TimestampedSlidingWindow,
+)
 from limits.typing import (
+    Any,
     Callable,
-    List,
     MemcachedClientP,
     Optional,
     P,
     R,
-    Tuple,
     Type,
     Union,
+    cast,
 )
 from limits.util import get_dependency
 
 
-class MemcachedStorage(Storage):
+class MemcachedStorage(Storage, SlidingWindowCounterSupport, TimestampedSlidingWindow):
     """
     Rate limit storage with memcached as backend.
 
@@ -70,7 +77,7 @@ class MemcachedStorage(Storage):
             options.pop("cluster_library", "pymemcache.client.hash")
         )
         self.client_getter = cast(
-            Callable[[ModuleType, List[Tuple[str, int]]], MemcachedClientP],
+            Callable[[ModuleType, list[tuple[str, int]]], MemcachedClientP],
             options.pop("client_getter", self.get_client),
         )
         self.options = options
@@ -86,11 +93,11 @@ class MemcachedStorage(Storage):
     @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.MemcacheError  # type: ignore[no-any-return]
 
     def get_client(
-        self, module: ModuleType, hosts: List[Tuple[str, int]], **kwargs: str
+        self, module: ModuleType, hosts: list[tuple[str, int]], **kwargs: str
     ) -> MemcachedClientP:
         """
         returns a memcached client.
@@ -142,8 +149,15 @@ class MemcachedStorage(Storage):
         """
         :param key: the key to get the counter value for
         """
+        return int(self.storage.get(key, "0"))
+
+    def get_many(self, keys: Iterable[str]) -> dict[str, Any]:  # type:ignore[explicit-any]
+        """
+        Return multiple counters at once
 
-        return int(self.storage.get(key) or 0)
+        :param keys: the keys to get the counter values for
+        """
+        return self.storage.get_many(keys)
 
     def clear(self, key: str) -> None:
         """
@@ -152,7 +166,12 @@ class MemcachedStorage(Storage):
         self.storage.delete(key)
 
     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,
+        set_expiration_key: bool = True,
     ) -> int:
         """
         increments the counter for a given rate limit key
@@ -162,41 +181,67 @@ class MemcachedStorage(Storage):
         :param elastic_expiry: whether to keep extending the rate limit
          window every hit.
         :param amount: the number to increment by
+        :param set_expiration_key: set the expiration key with the expiration time if needed. If set to False, the key will still expire, but memcached cannot provide the expiration time.
         """
-
-        if not self.call_memcached_func(
-            self.storage.add, key, amount, expiry, noreply=False
-        ):
-            value = self.storage.incr(key, amount) or amount
-
+        value = self.call_memcached_func(self.storage.incr, key, amount, noreply=False)
+        if value is not None:
             if elastic_expiry:
-                self.call_memcached_func(self.storage.touch, key, expiry)
-                self.call_memcached_func(
-                    self.storage.set,
-                    key + "/expires",
-                    expiry + time.time(),
-                    expire=expiry,
-                    noreply=False,
-                )
+                self.call_memcached_func(self.storage.touch, key, ceil(expiry))
+                if set_expiration_key:
+                    self.call_memcached_func(
+                        self.storage.set,
+                        self._expiration_key(key),
+                        expiry + time.time(),
+                        expire=ceil(expiry),
+                        noreply=False,
+                    )
 
             return value
         else:
-            self.call_memcached_func(
-                self.storage.set,
-                key + "/expires",
-                expiry + time.time(),
-                expire=expiry,
-                noreply=False,
-            )
-
-        return amount
+            if not self.call_memcached_func(
+                self.storage.add, key, amount, ceil(expiry), noreply=False
+            ):
+                value = self.storage.incr(key, amount) or amount
+
+                if elastic_expiry:
+                    self.call_memcached_func(self.storage.touch, key, ceil(expiry))
+                    if set_expiration_key:
+                        self.call_memcached_func(
+                            self.storage.set,
+                            self._expiration_key(key),
+                            expiry + time.time(),
+                            expire=ceil(expiry),
+                            noreply=False,
+                        )
+
+                return value
+            else:
+                if set_expiration_key:
+                    self.call_memcached_func(
+                        self.storage.set,
+                        self._expiration_key(key),
+                        expiry + time.time(),
+                        expire=ceil(expiry),
+                        noreply=False,
+                    )
+
+            return amount
 
     def get_expiry(self, key: str) -> float:
         """
         :param key: the key to get the expiry for
         """
 
-        return float(self.storage.get(key + "/expires") or time.time())
+        return float(self.storage.get(self._expiration_key(key)) or time.time())
+
+    def _expiration_key(self, key: str) -> str:
+        """
+        Return the expiration key for the given counter key.
+
+        Memcached doesn't natively return the expiration time or TTL for a given key,
+        so we implement the expiration time on a separate key.
+        """
+        return key + "/expires"
 
     def check(self) -> bool:
         """
@@ -212,3 +257,67 @@ class MemcachedStorage(Storage):
 
     def reset(self) -> Optional[int]:
         raise NotImplementedError
+
+    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, _ = self._get_sliding_window_info(
+            previous_key, current_key, expiry, now=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.
+            # We don't need the expiration key as it is estimated with the timestamps directly.
+            current_count = self.incr(
+                current_key, 2 * expiry, amount=amount, set_expiration_key=False
+            )
+            actualised_previous_ttl = min(0, previous_ttl - (time.time() - now))
+            weighted_count = (
+                previous_count * actualised_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.
+                self.call_memcached_func(
+                    self.storage.decr,
+                    current_key,
+                    amount,
+                    noreply=True,
+                )
+                return False
+            return True
+
+    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 self._get_sliding_window_info(previous_key, current_key, expiry, now)
+
+    def _get_sliding_window_info(
+        self, previous_key: str, current_key: str, expiry: int, now: float
+    ) -> tuple[int, float, int, float]:
+        result = self.get_many([previous_key, current_key])
+        previous_count, current_count = (
+            int(result.get(previous_key, 0)),
+            int(result.get(current_key, 0)),
+        )
+
+        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

limits/storage/memory.py

@@ -1,10 +1,18 @@
+from __future__ import annotations
+
 import threading
 import time
-from collections import Counter
+from collections import Counter, defaultdict
+from math import floor
 
 import limits.typing
-from limits.storage.base import MovingWindowSupport, Storage
-from limits.typing import Dict, List, Optional, Tuple, Type, Union
+from limits.storage.base import (
+    MovingWindowSupport,
+    SlidingWindowCounterSupport,
+    Storage,
+    TimestampedSlidingWindow,
+)
+from limits.typing import Optional, Type, Union
 
 
 class LockableEntry(threading._RLock):  # type: ignore
@@ -14,7 +22,9 @@ class LockableEntry(threading._RLock):  # type: ignore
         super().__init__()
 
 
-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,
@@ -28,8 +38,9 @@ class MemoryStorage(Storage, MovingWindowSupport):
         self, uri: Optional[str] = None, wrap_exceptions: bool = False, **_: str
     ):
         self.storage: limits.typing.Counter[str] = Counter()
-        self.expirations: Dict[str, float] = {}
-        self.events: Dict[str, List[LockableEntry]] = {}
+        self.locks: defaultdict[str, threading.RLock] = defaultdict(threading.RLock)
+        self.expirations: dict[str, float] = {}
+        self.events: dict[str, list[LockableEntry]] = {}
         self.timer = threading.Timer(0.01, self.__expire_events)
         self.timer.start()
         super().__init__(uri, wrap_exceptions=wrap_exceptions, **_)
@@ -37,7 +48,7 @@ class MemoryStorage(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 ValueError
 
     def __expire_events(self) -> None:
@@ -51,6 +62,7 @@ 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)
 
     def __schedule_expiry(self) -> None:
         if not self.timer.is_alive():
@@ -58,7 +70,7 @@ class MemoryStorage(Storage, MovingWindowSupport):
             self.timer.start()
 
     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
@@ -71,10 +83,25 @@ class MemoryStorage(Storage, MovingWindowSupport):
         """
         self.get(key)
         self.__schedule_expiry()
-        self.storage[key] += amount
+        with self.locks[key]:
+            self.storage[key] += amount
+
+            if elastic_expiry or self.storage[key] == amount:
+                self.expirations[key] = time.time() + expiry
+
+        return self.storage.get(key, 0)
+
+    def decr(self, key: str, amount: int = 1) -> int:
+        """
+        decrements the counter for a given rate limit key
 
-        if elastic_expiry or self.storage[key] == amount:
-            self.expirations[key] = time.time() + expiry
+        :param key: the key to decrement
+        :param amount: the number to decrement by
+        """
+        self.get(key)
+        self.__schedule_expiry()
+        with self.locks[key]:
+            self.storage[key] = max(self.storage[key] - amount, 0)
 
         return self.storage.get(key, 0)
 
@@ -86,6 +113,7 @@ class MemoryStorage(Storage, MovingWindowSupport):
         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)
 
@@ -96,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)
 
     def acquire_entry(self, key: str, limit: int, expiry: int, amount: int = 1) -> bool:
         """
@@ -143,7 +172,7 @@ class MemoryStorage(Storage, MovingWindowSupport):
             else 0
         )
 
-    def get_moving_window(self, key: str, limit: int, expiry: int) -> Tuple[float, int]:
+    def get_moving_window(self, key: str, limit: int, expiry: int) -> tuple[float, int]:
         """
         returns the starting point and the number of entries in the moving
         window
@@ -161,6 +190,63 @@ class MemoryStorage(Storage, MovingWindowSupport):
 
         return timestamp, acquired
 
+    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,
+            _,
+        ) = 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 = 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.
+                self.decr(current_key, amount)
+                # print("Concurrent call, reverting the counter increase")
+                return False
+            return True
+
+    def _get_sliding_window_info(
+        self,
+        previous_key: str,
+        current_key: str,
+        expiry: int,
+        now: float,
+    ) -> tuple[int, float, int, float]:
+        previous_count = self.get(previous_key)
+        current_count = 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
+
+    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 self._get_sliding_window_info(previous_key, current_key, expiry, now)
+
     def check(self) -> bool:
         """
         check if storage is healthy
@@ -173,4 +259,5 @@ class MemoryStorage(Storage, MovingWindowSupport):
         self.storage.clear()
         self.expirations.clear()
         self.events.clear()
+        self.locks.clear()
         return num_items