Update rate limiter

docs/util.rst

@@ -81,6 +81,7 @@ A simple rate limiter implementation which can be used in rules.
 The limiter is not rule bound so the same limiter can be used in multiples files.
 It also works as expected across rule reloads.
 
+
 Defining limits
 ^^^^^^^^^^^^^^^^^^
 Limits can either be explicitly added or through a textual description.
@@ -100,10 +101,10 @@ Examples:
 * ``300 / hour``
 
 
-Elastic expiry
-^^^^^^^^^^^^^^^^^^
+Fixed window elastic expiry algorithm
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-The rate limiter implements a fixed window with elastic expiry.
+This algorithm implements a fixed window with elastic expiry.
 That means if the limit is hit the interval time will be increased by the expiry time.
 
 For example ``3 per minute``:
@@ -118,6 +119,14 @@ For example ``3 per minute``:
   will also get rejected and the intervall now goes from ``00:00:00`` - ``00:02:30``.
 
 
+Leaky bucket algorithm
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The leaky bucket algorithm is based on the analogy of a bucket that leaks at a constant rate.
+As long as the bucket is not full the hits will pass. If the bucket overflows the hits will get rejected.
+Since the bucket leaks at a constant rate it will gradually get empty again thus allowing hits to pass again.
+
+
 Example
 ^^^^^^^^^^^^^^^^^^
 
@@ -128,10 +137,12 @@ Example
     # Create or get existing, name is case insensitive
     limiter = RateLimiter('MyRateLimiterName')
 
-    # define limits, duplicate limits will only be added once
+    # define limits, duplicate limits of the same algorithm will only be added once
     limiter.add_limit(5, 60)   # add limits explicitly
     limiter.parse_limits('5 per minute').parse_limits('5 in 60s', '5/60seconds')  # add limits through text
 
+    # add additional limit with leaky bucket algorithm
+    limiter.add_limit(10, 120, algorithm='leaky_bucket')
 
     # Test the limit without increasing the hits
     for _ in range(100):
@@ -155,6 +166,19 @@ Documentation
 
 .. autoclass:: HABApp.util.rate_limiter.limiter.Limiter
    :members:
+   :inherited-members:
+
+.. autoclass:: HABApp.util.rate_limiter.limiter.LimiterInfo
+   :members:
+   :inherited-members:
+
+.. autoclass:: HABApp.util.rate_limiter.limiter.FixedWindowElasticExpiryLimitInfo
+   :members:
+   :inherited-members:
+
+.. autoclass:: HABApp.util.rate_limiter.limiter.LeakyBucketLimitInfo
+   :members:
+   :inherited-members:
 
 
 Statistics

src/HABApp/util/functions/min_max.py

@@ -2,8 +2,9 @@
 from builtins import min as _min
 
 
-def max(*args, default=None):
-    """Behaves like the built in max function but ignores any ``None`` values. e.g. ``max([1, None, 2]) == 2``.
+# noinspection PyShadowingBuiltins
+def max(*args, default=None):   # noqa: A001
+    """Behaves like the built-in max function but ignores any ``None`` values. e.g. ``max([1, None, 2]) == 2``.
     If the iterable is empty ``default`` will be returned.
 
     :param args: Single iterable or 1..n arguments
@@ -16,8 +17,9 @@ def max(*args, default=None):
     )
 
 
-def min(*args, default=None):
-    """Behaves like the built in min function but ignores any ``None`` values. e.g. ``min([1, None, 2]) == 1``.
+# noinspection PyShadowingBuiltins
+def min(*args, default=None):   # noqa: A001
+    """Behaves like the built-in min function but ignores any ``None`` values. e.g. ``min([1, None, 2]) == 1``.
     If the iterable is empty ``default`` will be returned.
 
     :param args: Single iterable or 1..n arguments

src/HABApp/util/rate_limiter/limiter.py

@@ -1,82 +1,84 @@
-import re
 from dataclasses import dataclass
-from time import monotonic
-from typing import Final, List, Tuple
+from typing import Final, List, Literal, Tuple, TypeAlias, Union
 
-from .rate_limit import RateLimit, RateLimitInfo
+from HABApp.core.const.const import StrEnum
 
-
-LIMIT_REGEX = re.compile(
-    r"""
-    \s* ([1-9][0-9]*)
-    \s* (/|per|in)
-    \s* ([1-9][0-9]*)?
-    \s* (s|sec|second|m|min|minute|h|hour|day|month|year)s?
-    \s*""",
-    re.IGNORECASE | re.VERBOSE,
+from .limits import (
+    BaseRateLimit,
+    FixedWindowElasticExpiryLimit,
+    FixedWindowElasticExpiryLimitInfo,
+    LeakyBucketLimit,
+    LeakyBucketLimitInfo,
 )
+from .parser import parse_limit
 
 
-def parse_limit(text: str) -> Tuple[int, int]:
-    if not isinstance(text, str) or not (m := LIMIT_REGEX.fullmatch(text)):
-        msg = f'Invalid limit string: "{text:s}"'
-        raise ValueError(msg)
-
-    count, per, factor, interval = m.groups()
+class LimitTypeEnum(StrEnum):
+    LEAKY_BUCKET = 'leaky_bucket'
+    FIXED_WINDOW_ELASTIC_EXPIRY = 'fixed_window_elastic_expiry'
 
-    interval_secs = {
-        's': 1, 'sec': 1, 'second': 1, 'm': 60, 'min': 60, 'minute': 60, 'hour': 3600, 'h': 3600,
-        'day': 24 * 3600, 'month': 30 * 24 * 3600, 'year': 365 * 24 * 3600
-    }[interval]
 
-    return int(count), int(1 if factor is None else factor) * interval_secs
+LIMITER_ALGORITHM_HINT: TypeAlias = Literal[LimitTypeEnum.LEAKY_BUCKET, LimitTypeEnum.FIXED_WINDOW_ELASTIC_EXPIRY]
 
 
 class Limiter:
     def __init__(self, name: str):
         self._name: Final = name
-        self._limits: Tuple[RateLimit, ...] = ()
-        self._skipped = 0
+        self._limits: Tuple[BaseRateLimit, ...] = ()
+        self._skips = 0
 
     def __repr__(self):
         return f'<{self.__class__.__name__} {self._name:s}>'
 
-    def add_limit(self, allowed: int, expiry: int) -> 'Limiter':
+    def add_limit(self, allowed: int, interval: int,
+                  algorithm: LIMITER_ALGORITHM_HINT = LimitTypeEnum.FIXED_WINDOW_ELASTIC_EXPIRY) -> 'Limiter':
         """Add a new rate limit
 
         :param allowed: How many hits are allowed
-        :param expiry: Interval in seconds
+        :param interval: Interval in seconds
+        :param algorithm: Which algorithm should this limit use
         """
         if allowed <= 0 or not isinstance(allowed, int):
             msg = f'Allowed must be an int >= 0, is {allowed} ({type(allowed)})'
             raise ValueError(msg)
 
-        if expiry <= 0 or not isinstance(expiry, int):
-            msg = f'Expire time must be an int >= 0, is {expiry} ({type(expiry)})'
+        if interval <= 0 or not isinstance(interval, int):
+            msg = f'Expire time must be an int >= 0, is {interval} ({type(interval)})'
             raise ValueError(msg)
 
+        algo = LimitTypeEnum(algorithm)
+        if algo is LimitTypeEnum.FIXED_WINDOW_ELASTIC_EXPIRY:
+            cls = FixedWindowElasticExpiryLimit
+        elif algo is LimitTypeEnum.LEAKY_BUCKET:
+            cls = LeakyBucketLimit
+        else:
+            raise ValueError()
+
+        # Check if we have already added an algorithm with these parameters
         for window in self._limits:
-            if window.allowed == allowed and window.expiry == expiry:
+            if isinstance(window, cls) and window.allowed == allowed and window.interval == interval:
                 return self
 
-        limit = RateLimit(allowed, expiry)
-        self._limits = tuple(sorted([*self._limits, limit], key=lambda x: x.expiry))
+        limit = cls(allowed, interval)
+        self._limits = tuple(sorted([*self._limits, limit], key=lambda x: x.interval))
         return self
 
-    def parse_limits(self, *text: str) -> 'Limiter':
+    def parse_limits(self, *text: str,
+                     algorithm: LIMITER_ALGORITHM_HINT = LimitTypeEnum.FIXED_WINDOW_ELASTIC_EXPIRY) -> 'Limiter':
         """Add one or more limits in textual form, e.g. ``5 in 60s``, ``10 per hour`` or ``10/15 mins``.
         If the limit does already exist it will not be added again.
 
         :param text: textual description of limit
+        :param algorithm: Which algorithm should these limits use
         """
         for limit in [parse_limit(t) for t in text]:
-            self.add_limit(*limit)
+            self.add_limit(*limit, algorithm=algorithm)
         return self
 
     def allow(self) -> bool:
         """Test the limit.
 
-        :return: True if allowed, False if forbidden
+        :return: ``True`` if allowed, ``False`` if forbidden
         """
         allow = True
         clear_skipped = True
@@ -93,17 +95,17 @@ def allow(self) -> bool:
                 clear_skipped = False
 
         if clear_skipped:
-            self._skipped = 0
+            self._skips = 0
 
         if not allow:
-            self._skipped += 1
+            self._skips += 1
 
         return allow
 
     def test_allow(self) -> bool:
         """Test the limit without hitting it. Calling this will not increase the hit counter.
 
-        :return: True if allowed, False if forbidden
+        :return: ``True`` if allowed, ``False`` if forbidden
         """
         allow = True
         clear_skipped = True
@@ -119,25 +121,20 @@ def test_allow(self) -> bool:
                 clear_skipped = False
 
         if clear_skipped:
-            self._skipped = 0
+            self._skips = 0
         return allow
 
     def info(self) -> 'LimiterInfo':
         """Get some info about the limiter and the defined windows
         """
-        now = monotonic()
-        remaining = max((w.stop for w in self._limits if w.hits), default=now) - now
-        if remaining <= 0:
-            remaining = 0
 
         return LimiterInfo(
-            time_remaining=remaining, skipped=self._skipped,
-            limits=[limit.window_info() for limit in self._limits]
+            skips=self._skips,
+            limits=[limit.info() for limit in self._limits]
         )
 
 
 @dataclass
 class LimiterInfo:
-    time_remaining: float   #: time remaining until skipped will reset
-    skipped: int            #: how many entries were skipped
-    limits: List['RateLimitInfo']    # Info for every window
+    skips: int            #: How many entries were skipped
+    limits: List[Union[FixedWindowElasticExpiryLimitInfo, LeakyBucketLimitInfo]]    #: Info for every limit

src/HABApp/util/rate_limiter/limits/__init__.py

@@ -0,0 +1,3 @@
+from .base import BaseRateLimit
+from .fixed_window import FixedWindowElasticExpiryLimit, FixedWindowElasticExpiryLimitInfo
+from .leaky_bucket import LeakyBucketLimit, LeakyBucketLimitInfo

src/HABApp/util/rate_limiter/limits/base.py

@@ -0,0 +1,73 @@
+from dataclasses import dataclass
+from typing import Final
+
+
+@dataclass
+class BaseRateLimitInfo:
+    hits: int               #: Hits
+    skips: int              #: Skips
+    limit: int              #: Boundary
+
+    @property
+    def hits_remaining(self) -> int:
+        return self.limit - self.hits
+
+
+class BaseRateLimit:
+    def __init__(self, allowed: int, interval: int):
+        super().__init__()
+        assert allowed > 0, allowed
+        assert interval > 0, interval
+
+        self.interval: Final = interval
+        self.allowed: Final = allowed
+
+        self.hits: int = 0
+        self.skips: int = 0
+
+    def repr_text(self) -> str:
+        return ''
+
+    def __repr__(self):
+        return (
+            f'<{self.__class__.__name__} hits={self.hits:d}/{self.allowed:d} interval={self.interval:d}s '
+            f'{self.repr_text():s}>'
+        )
+
+    def do_test_allow(self):
+        raise NotImplementedError()
+
+    def do_allow(self):
+        raise NotImplementedError()
+
+    def do_deny(self):
+        raise NotImplementedError()
+
+    def info(self) -> BaseRateLimitInfo:
+        raise NotImplementedError()
+
+    def allow(self, weight: int = 1) -> bool:
+        if not isinstance(weight, int) or weight <= 0:
+            msg = f'weight must be an int > 0, is {weight}'
+            raise ValueError(msg)
+
+        self.do_allow()
+
+        self.hits += weight
+        if self.hits <= self.allowed:
+            return True
+
+        self.skips += 1
+        self.hits = self.allowed
+
+        if self.do_deny:
+            self.do_deny()
+        return False
+
+    def test_allow(self, weight: int = 1) -> bool:
+        if not isinstance(weight, int) or weight <= 0:
+            msg = f'weight must be an int > 0, is {weight}'
+            raise ValueError(msg)
+
+        self.do_test_allow()
+        return self.hits + weight <= self.allowed

src/HABApp/util/rate_limiter/limits/fixed_window.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass
+from time import monotonic
+
+from .base import BaseRateLimit, BaseRateLimitInfo
+
+
+@dataclass
+class FixedWindowElasticExpiryLimitInfo(BaseRateLimitInfo):
+    time_remaining: float   #: Time remaining until this window will reset
+
+
+class FixedWindowElasticExpiryLimit(BaseRateLimit):
+    def __init__(self, allowed: int, interval: int):
+        super().__init__(allowed, interval)
+
+        self.start: float = -1.0
+        self.stop: float = -1.0
+
+    def repr_text(self):
+        return f'window={self.stop - self.start:.0f}s'
+
+    def do_test_allow(self):
+        if self.stop <= monotonic():
+            self.hits = 0
+            self.skips = 0
+
+    def do_allow(self):
+        now = monotonic()
+
+        if self.stop <= now:
+            self.hits = 0
+            self.skips = 0
+            self.start = now
+            self.stop = now + self.interval
+
+    def do_deny(self):
+        self.stop = monotonic() + self.interval
+
+    def info(self) -> FixedWindowElasticExpiryLimitInfo:
+        self.do_test_allow()
+
+        remaining = self.stop - monotonic()
+        if remaining <= 0:
+            remaining = 0
+
+        if not remaining and not self.hits:
+            remaining = self.interval
+
+        return FixedWindowElasticExpiryLimitInfo(
+            time_remaining=remaining, hits=self.hits, skips=self.skips, limit=self.allowed
+        )