vutran1710 / PyrateLimiter

⚔️Python Rate-Limiter using Leaky-Bucket Algorithm Family
https://pyratelimiter.readthedocs.io
MIT License
334 stars 36 forks source link

async BucketFactory #133

Closed grebelsm closed 11 months ago

grebelsm commented 11 months ago

I am trying to implement my own BucketFactory with custom routing logic to be used with a Redis backend asynchronously. When implementing my own BucketFactory, can I redefine get() to be async like so?

class MyBucketFactory(BucketFactory):
    def __init__(self):
        self.base_clock = TimeClock()
        pool = ConnectionPool.from_url("redis://localhost:6379")
        self.redis_db = Redis(connection_pool=pool)
        self.rate = Rate(10, Duration.SECOND)

    def wrap_item(self, request: str, weight: int = 1) -> RateItem:
        """Time-stamping item, return a RateItem"""
        now = self.base_clock.now()
        return RateItem(f"{request.x}_{request.y}", now, weight=weight)

    async def get(self, _item: RateItem) -> AbstractBucket:
        bucket = await RedisBucket.init([self.rate], self.redis_db, _item.name)
        bucket.leak(self.base_clock.now())
        return bucket

then

limiter = Limiter(MyBucketFactory())

then

await limiter.try_acquire(request) ?

vutran1710 commented 11 months ago

short answer is no - but i believe you don't actually need it.

the redis bucket init method only needs to be async once for the very first bucket. - as it needs to load a lua script into redis for the algorithm to run.

for the later buckets, if you want to create them dynamically during routing, you can directly call the constructor and passing the existing lua-scripts hash to it

for example:


script_hash = your_first_ever_redis_bucket.script_hash

new_bucket = RedisBucket(rates, db, bucket_key, script_hash)


Sorry, im on my phone so cannot put on a proper snippet format.

vutran1710 commented 11 months ago

or, you can actively load the script first before creating any bucket, then pass it to the bucket constructor. with that you wont need async-get on the bucket-factory

grebelsm commented 11 months ago

Thanks for the quick response, so basically something like this?

class MyBucketFactory(BucketFactory):
    def __init__(self):
        pool = ConnectionPool.from_url("redis://localhost:6379")
        self.redis_db = AsyncRedis(connection_pool=pool)
        self.script_hash = None
        self.rate = Rate(10, Duration.SECOND)
        self.base_clock = TimeClock()

    async def _init_script(self):
        self.script_hash = await self.redis_db.script_load(LuaScript.PUT_ITEM)

    def wrap_item(self, request: str, weight: int = 1) -> RateItem:
        """Time-stamping item, return a RateItem"""
        now = self.base_clock.now()
        return RateItem(f"{request.x}_{request.y}", now, weight=weight)

    def get(self, _item: RateItem) -> AbstractBucket:
        bucket = RedisBucket([self.rate], self.redis_db, _item.name, self.script_hash)
        bucket.leak(base_clock.now())
        return bucket

then

    bucket_factory = MyBucketFactory()
    await bucket_factory._init_script()
    limiter = Limiter(bucket_factory)

    res = await limiter.try_acquire(request_1)

?

and line bucket = RedisBucket([self.rate], self.redis_db, _item.name, self.script_hash) will get existing bucket based on _item.name, or create a new bucket, both using async operations?

grebelsm commented 11 months ago

Also, regarding bucket.leak(base_clock.now()), I want to leak the bucket before I return it, but Im guessing this needs to be awaited. Is there a better way for me to do that?

vutran1710 commented 11 months ago

Also, regarding bucket.leak(base_clock.now()), I want to leak the bucket before I return it, but Im guessing this needs to be awaited. Is there a better way for me to do that?

I did this with pyrate-limiter v2 then realized that it is not a good way to leak, since it adds extra time to consuming step and leads to timestampt incorrectness, so i suggest you to avoid it.

the schedule-leak method of BucketFactory base class made for just that and you dont have to do anything except every time you create a bucket, remeber to pass it like this


bucket = RedisBucket(...)

self.schedule_leak(bucket, self.base_clock)


vutran1710 commented 11 months ago

you can look at the Leaking section in the documentation for more details

grebelsm commented 11 months ago

great thanks for the help, closing it out