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

limits/storage/memcached.py

@@ -2,26 +2,31 @@ 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 +75,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 +91,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 +147,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 +164,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 +179,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 +255,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,16 @@
 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 +20,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 +36,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 +46,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 +60,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 +68,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 +81,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
+            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
+
+        :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 +111,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 +122,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 +170,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 +188,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 +257,5 @@ class MemoryStorage(Storage, MovingWindowSupport):
         self.storage.clear()
         self.expirations.clear()
         self.events.clear()
+        self.locks.clear()
         return num_items