Initial commit
This commit is contained in:
kdusek
@@ -0,0 +1,4 @@
|
||||
from .bucket import * # noqa
|
||||
from .clock import * # noqa
|
||||
from .rate import * # noqa
|
||||
from .wrappers import * # noqa
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,304 @@
|
||||
""" Implement this class to create
|
||||
a workable bucket for Limiter to use
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
from abc import ABC
|
||||
from abc import abstractmethod
|
||||
from collections import defaultdict
|
||||
from inspect import isawaitable
|
||||
from inspect import iscoroutine
|
||||
from threading import Thread
|
||||
from typing import Awaitable
|
||||
from typing import Dict
|
||||
from typing import List
|
||||
from typing import Optional
|
||||
from typing import Type
|
||||
from typing import Union
|
||||
|
||||
from .clock import AbstractClock
|
||||
from .rate import Rate
|
||||
from .rate import RateItem
|
||||
|
||||
logger = logging.getLogger("pyrate_limiter")
|
||||
|
||||
|
||||
class AbstractBucket(ABC):
|
||||
"""Base bucket interface
|
||||
Assumption: len(rates) always > 0
|
||||
TODO: allow empty rates
|
||||
"""
|
||||
|
||||
rates: List[Rate]
|
||||
failing_rate: Optional[Rate] = None
|
||||
|
||||
@abstractmethod
|
||||
def put(self, item: RateItem) -> Union[bool, Awaitable[bool]]:
|
||||
"""Put an item (typically the current time) in the bucket
|
||||
return true if successful, otherwise false
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def leak(
|
||||
self,
|
||||
current_timestamp: Optional[int] = None,
|
||||
) -> Union[int, Awaitable[int]]:
|
||||
"""leaking bucket - removing items that are outdated"""
|
||||
|
||||
@abstractmethod
|
||||
def flush(self) -> Union[None, Awaitable[None]]:
|
||||
"""Flush the whole bucket
|
||||
- Must remove `failing-rate` after flushing
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def count(self) -> Union[int, Awaitable[int]]:
|
||||
"""Count number of items in the bucket"""
|
||||
|
||||
@abstractmethod
|
||||
def peek(self, index: int) -> Union[Optional[RateItem], Awaitable[Optional[RateItem]]]:
|
||||
"""Peek at the rate-item at a specific index in latest-to-earliest order
|
||||
NOTE: The reason we cannot peek from the start of the queue(earliest-to-latest) is
|
||||
we can't really tell how many outdated items are still in the queue
|
||||
"""
|
||||
|
||||
def waiting(self, item: RateItem) -> Union[int, Awaitable[int]]:
|
||||
"""Calculate time until bucket become availabe to consume an item again"""
|
||||
if self.failing_rate is None:
|
||||
return 0
|
||||
|
||||
assert item.weight > 0, "Item's weight must > 0"
|
||||
|
||||
if item.weight > self.failing_rate.limit:
|
||||
return -1
|
||||
|
||||
bound_item = self.peek(self.failing_rate.limit - item.weight)
|
||||
|
||||
if bound_item is None:
|
||||
# NOTE: No waiting, bucket is immediately ready
|
||||
return 0
|
||||
|
||||
def _calc_waiting(inner_bound_item: RateItem) -> int:
|
||||
assert self.failing_rate is not None # NOTE: silence mypy
|
||||
lower_time_bound = item.timestamp - self.failing_rate.interval
|
||||
upper_time_bound = inner_bound_item.timestamp
|
||||
return upper_time_bound - lower_time_bound
|
||||
|
||||
async def _calc_waiting_async() -> int:
|
||||
nonlocal bound_item
|
||||
|
||||
while isawaitable(bound_item):
|
||||
bound_item = await bound_item
|
||||
|
||||
if bound_item is None:
|
||||
# NOTE: No waiting, bucket is immediately ready
|
||||
return 0
|
||||
|
||||
assert isinstance(bound_item, RateItem)
|
||||
return _calc_waiting(bound_item)
|
||||
|
||||
if isawaitable(bound_item):
|
||||
return _calc_waiting_async()
|
||||
|
||||
assert isinstance(bound_item, RateItem)
|
||||
return _calc_waiting(bound_item)
|
||||
|
||||
def limiter_lock(self) -> Optional[object]: # type: ignore
|
||||
"""An additional lock to be used by Limiter in-front of the thread lock.
|
||||
Intended for multiprocessing environments where a thread lock is insufficient.
|
||||
"""
|
||||
return None
|
||||
|
||||
|
||||
class Leaker(Thread):
|
||||
"""Responsible for scheduling buckets' leaking at the background either
|
||||
through a daemon task(for sync buckets) or a task using asyncio.Task
|
||||
"""
|
||||
|
||||
daemon = True
|
||||
name = "PyrateLimiter's Leaker"
|
||||
sync_buckets: Optional[Dict[int, AbstractBucket]] = None
|
||||
async_buckets: Optional[Dict[int, AbstractBucket]] = None
|
||||
clocks: Optional[Dict[int, AbstractClock]] = None
|
||||
leak_interval: int = 10_000
|
||||
aio_leak_task: Optional[asyncio.Task] = None
|
||||
|
||||
def __init__(self, leak_interval: int):
|
||||
self.sync_buckets = defaultdict()
|
||||
self.async_buckets = defaultdict()
|
||||
self.clocks = defaultdict()
|
||||
self.leak_interval = leak_interval
|
||||
super().__init__()
|
||||
|
||||
def register(self, bucket: AbstractBucket, clock: AbstractClock):
|
||||
"""Register a new bucket with its associated clock"""
|
||||
assert self.sync_buckets is not None
|
||||
assert self.clocks is not None
|
||||
assert self.async_buckets is not None
|
||||
|
||||
try_leak = bucket.leak(0)
|
||||
bucket_id = id(bucket)
|
||||
|
||||
if iscoroutine(try_leak):
|
||||
try_leak.close()
|
||||
self.async_buckets[bucket_id] = bucket
|
||||
else:
|
||||
self.sync_buckets[bucket_id] = bucket
|
||||
|
||||
self.clocks[bucket_id] = clock
|
||||
|
||||
def deregister(self, bucket_id: int) -> bool:
|
||||
"""Deregister a bucket"""
|
||||
if self.sync_buckets and bucket_id in self.sync_buckets:
|
||||
del self.sync_buckets[bucket_id]
|
||||
assert self.clocks
|
||||
del self.clocks[bucket_id]
|
||||
return True
|
||||
|
||||
if self.async_buckets and bucket_id in self.async_buckets:
|
||||
del self.async_buckets[bucket_id]
|
||||
assert self.clocks
|
||||
del self.clocks[bucket_id]
|
||||
|
||||
if not self.async_buckets and self.aio_leak_task:
|
||||
self.aio_leak_task.cancel()
|
||||
self.aio_leak_task = None
|
||||
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
async def _leak(self, buckets: Dict[int, AbstractBucket]) -> None:
|
||||
assert self.clocks
|
||||
|
||||
while buckets:
|
||||
try:
|
||||
for bucket_id, bucket in list(buckets.items()):
|
||||
clock = self.clocks[bucket_id]
|
||||
now = clock.now()
|
||||
|
||||
while isawaitable(now):
|
||||
now = await now
|
||||
|
||||
assert isinstance(now, int)
|
||||
leak = bucket.leak(now)
|
||||
|
||||
while isawaitable(leak):
|
||||
leak = await leak
|
||||
|
||||
assert isinstance(leak, int)
|
||||
|
||||
await asyncio.sleep(self.leak_interval / 1000)
|
||||
except RuntimeError as e:
|
||||
logger.info("Leak task stopped due to event loop shutdown. %s", e)
|
||||
return
|
||||
|
||||
def leak_async(self):
|
||||
if self.async_buckets and not self.aio_leak_task:
|
||||
self.aio_leak_task = asyncio.create_task(self._leak(self.async_buckets))
|
||||
|
||||
def run(self) -> None:
|
||||
""" Override the original method of Thread
|
||||
Not meant to be called directly
|
||||
"""
|
||||
assert self.sync_buckets
|
||||
asyncio.run(self._leak(self.sync_buckets))
|
||||
|
||||
def start(self) -> None:
|
||||
""" Override the original method of Thread
|
||||
Call to run leaking sync buckets
|
||||
"""
|
||||
if self.sync_buckets and not self.is_alive():
|
||||
super().start()
|
||||
|
||||
|
||||
class BucketFactory(ABC):
|
||||
"""Asbtract BucketFactory class.
|
||||
It is reserved for user to implement/override this class with
|
||||
his own bucket-routing/creating logic
|
||||
"""
|
||||
|
||||
_leaker: Optional[Leaker] = None
|
||||
_leak_interval: int = 10_000
|
||||
|
||||
@property
|
||||
def leak_interval(self) -> int:
|
||||
"""Retrieve leak-interval from inner Leaker task"""
|
||||
if not self._leaker:
|
||||
return self._leak_interval
|
||||
return self._leaker.leak_interval
|
||||
|
||||
@leak_interval.setter
|
||||
def leak_interval(self, value: int):
|
||||
"""Set leak-interval for inner Leaker task"""
|
||||
if self._leaker:
|
||||
self._leaker.leak_interval = value
|
||||
self._leak_interval = value
|
||||
|
||||
@abstractmethod
|
||||
def wrap_item(
|
||||
self,
|
||||
name: str,
|
||||
weight: int = 1,
|
||||
) -> Union[RateItem, Awaitable[RateItem]]:
|
||||
"""Add the current timestamp to the receiving item using any clock backend
|
||||
- Turn it into a RateItem
|
||||
- Can return either a coroutine or a RateItem instance
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get(self, item: RateItem) -> Union[AbstractBucket, Awaitable[AbstractBucket]]:
|
||||
"""Get the corresponding bucket to this item"""
|
||||
|
||||
def create(
|
||||
self,
|
||||
clock: AbstractClock,
|
||||
bucket_class: Type[AbstractBucket],
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> AbstractBucket:
|
||||
"""Creating a bucket dynamically"""
|
||||
bucket = bucket_class(*args, **kwargs)
|
||||
self.schedule_leak(bucket, clock)
|
||||
return bucket
|
||||
|
||||
def schedule_leak(self, new_bucket: AbstractBucket, associated_clock: AbstractClock) -> None:
|
||||
"""Schedule all the buckets' leak, reset bucket's failing rate"""
|
||||
assert new_bucket.rates, "Bucket rates are not set"
|
||||
|
||||
if not self._leaker:
|
||||
self._leaker = Leaker(self.leak_interval)
|
||||
|
||||
self._leaker.register(new_bucket, associated_clock)
|
||||
self._leaker.start()
|
||||
self._leaker.leak_async()
|
||||
|
||||
def get_buckets(self) -> List[AbstractBucket]:
|
||||
"""Iterator over all buckets in the factory
|
||||
"""
|
||||
if not self._leaker:
|
||||
return []
|
||||
|
||||
buckets = []
|
||||
|
||||
if self._leaker.sync_buckets:
|
||||
for _, bucket in self._leaker.sync_buckets.items():
|
||||
buckets.append(bucket)
|
||||
|
||||
if self._leaker.async_buckets:
|
||||
for _, bucket in self._leaker.async_buckets.items():
|
||||
buckets.append(bucket)
|
||||
|
||||
return buckets
|
||||
|
||||
def dispose(self, bucket: Union[int, AbstractBucket]) -> bool:
|
||||
"""Delete a bucket from the factory"""
|
||||
if isinstance(bucket, AbstractBucket):
|
||||
bucket = id(bucket)
|
||||
|
||||
assert isinstance(bucket, int), "not valid bucket id"
|
||||
|
||||
if not self._leaker:
|
||||
return False
|
||||
|
||||
return self._leaker.deregister(bucket)
|
||||
@@ -0,0 +1,12 @@
|
||||
from abc import ABC
|
||||
from abc import abstractmethod
|
||||
from typing import Awaitable
|
||||
from typing import Union
|
||||
|
||||
|
||||
class AbstractClock(ABC):
|
||||
"""Clock that return timestamp for `now`"""
|
||||
|
||||
@abstractmethod
|
||||
def now(self) -> Union[int, Awaitable[int]]:
|
||||
"""Get time as of now, in miliseconds"""
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Unit classes that deals with rate, item & duration
|
||||
"""
|
||||
from enum import Enum
|
||||
from typing import Union
|
||||
|
||||
|
||||
class Duration(Enum):
|
||||
"""Interval helper class"""
|
||||
|
||||
SECOND = 1000
|
||||
MINUTE = 1000 * 60
|
||||
HOUR = 1000 * 60 * 60
|
||||
DAY = 1000 * 60 * 60 * 24
|
||||
WEEK = 1000 * 60 * 60 * 24 * 7
|
||||
|
||||
def __mul__(self, mutiplier: float) -> int:
|
||||
return int(self.value * mutiplier)
|
||||
|
||||
def __rmul__(self, multiplier: float) -> int:
|
||||
return self.__mul__(multiplier)
|
||||
|
||||
def __add__(self, another_duration: Union["Duration", int]) -> int:
|
||||
return self.value + int(another_duration)
|
||||
|
||||
def __radd__(self, another_duration: Union["Duration", int]) -> int:
|
||||
return self.__add__(another_duration)
|
||||
|
||||
def __int__(self) -> int:
|
||||
return self.value
|
||||
|
||||
def __eq__(self, duration: object) -> bool:
|
||||
if not isinstance(duration, (Duration, int)):
|
||||
return NotImplemented
|
||||
|
||||
return self.value == int(duration)
|
||||
|
||||
@staticmethod
|
||||
def readable(value: int) -> str:
|
||||
notes = [
|
||||
(Duration.WEEK, "w"),
|
||||
(Duration.DAY, "d"),
|
||||
(Duration.HOUR, "h"),
|
||||
(Duration.MINUTE, "m"),
|
||||
(Duration.SECOND, "s"),
|
||||
]
|
||||
|
||||
for note, shorten in notes:
|
||||
if value >= note.value:
|
||||
decimal_value = value / note.value
|
||||
return f"{decimal_value:0.1f}{shorten}" # noqa: E231
|
||||
|
||||
return f"{value}ms"
|
||||
|
||||
|
||||
class RateItem:
|
||||
"""RateItem is a wrapper for bucket to work with"""
|
||||
|
||||
name: str
|
||||
weight: int
|
||||
timestamp: int
|
||||
|
||||
def __init__(self, name: str, timestamp: int, weight: int = 1):
|
||||
self.name = name
|
||||
self.timestamp = timestamp
|
||||
self.weight = weight
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"RateItem(name={self.name}, weight={self.weight}, timestamp={self.timestamp})"
|
||||
|
||||
|
||||
class Rate:
|
||||
"""Rate definition.
|
||||
|
||||
Args:
|
||||
limit: Number of requests allowed within ``interval``
|
||||
interval: Time interval, in miliseconds
|
||||
"""
|
||||
|
||||
limit: int
|
||||
interval: int
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
limit: int,
|
||||
interval: Union[int, Duration],
|
||||
):
|
||||
self.limit = limit
|
||||
self.interval = int(interval)
|
||||
assert self.interval
|
||||
assert self.limit
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"limit={self.limit}/{Duration.readable(self.interval)}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"limit={self.limit}/{self.interval}"
|
||||
@@ -0,0 +1,77 @@
|
||||
""" Wrappers over different abstract types
|
||||
"""
|
||||
from inspect import isawaitable
|
||||
from typing import Optional
|
||||
|
||||
from .bucket import AbstractBucket
|
||||
from .rate import RateItem
|
||||
|
||||
|
||||
class BucketAsyncWrapper(AbstractBucket):
|
||||
"""BucketAsyncWrapper is a wrapping over any bucket
|
||||
that turns a async/synchronous bucket into an async one
|
||||
"""
|
||||
|
||||
def __init__(self, bucket: AbstractBucket):
|
||||
assert isinstance(bucket, AbstractBucket)
|
||||
self.bucket = bucket
|
||||
|
||||
async def put(self, item: RateItem):
|
||||
result = self.bucket.put(item)
|
||||
|
||||
while isawaitable(result):
|
||||
result = await result
|
||||
|
||||
return result
|
||||
|
||||
async def count(self):
|
||||
result = self.bucket.count()
|
||||
|
||||
while isawaitable(result):
|
||||
result = await result
|
||||
|
||||
return result
|
||||
|
||||
async def leak(self, current_timestamp: Optional[int] = None) -> int:
|
||||
result = self.bucket.leak(current_timestamp)
|
||||
|
||||
while isawaitable(result):
|
||||
result = await result
|
||||
|
||||
assert isinstance(result, int)
|
||||
return result
|
||||
|
||||
async def flush(self) -> None:
|
||||
result = self.bucket.flush()
|
||||
|
||||
while isawaitable(result):
|
||||
# TODO: AbstractBucket.flush() may not have correct type annotation?
|
||||
result = await result # type: ignore
|
||||
|
||||
return None
|
||||
|
||||
async def peek(self, index: int) -> Optional[RateItem]:
|
||||
item = self.bucket.peek(index)
|
||||
|
||||
while isawaitable(item):
|
||||
item = await item
|
||||
|
||||
assert item is None or isinstance(item, RateItem)
|
||||
return item
|
||||
|
||||
async def waiting(self, item: RateItem) -> int:
|
||||
wait = super().waiting(item)
|
||||
|
||||
if isawaitable(wait):
|
||||
wait = await wait
|
||||
|
||||
assert isinstance(wait, int)
|
||||
return wait
|
||||
|
||||
@property
|
||||
def failing_rate(self):
|
||||
return self.bucket.failing_rate
|
||||
|
||||
@property
|
||||
def rates(self):
|
||||
return self.bucket.rates
|
||||
Reference in New Issue
Block a user