throttler fix
This commit is contained in:
1 parent
535826fce6
commit
3e21b904a2
1 file changed
+16
-12
+16
-12
@@ -1,33 +1,37 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
from asyncio_throttle import Throttler as AsyncThrottler
|
from asyncio_throttle import Throttler as AsyncThrottler
|
||||||
|
|
||||||
|
|
||||||
class GlobalThrottler:
|
class GlobalThrottler:
|
||||||
_instance = None
|
_instance = None
|
||||||
_lock = asyncio.Lock()
|
_lock = asyncio.Lock()
|
||||||
|
|
||||||
def __init__(self, rate_limit: int = 500, period: int = 60) -> None:
|
def __init__(self, rate_limit: int = 500, period: int = 60) -> None:
|
||||||
if not hasattr(self, '_initialized'):
|
if not hasattr(self, "_initialized"):
|
||||||
self.rate_limit = rate_limit
|
self.rate_limit = rate_limit
|
||||||
self.period = period
|
self.period = period
|
||||||
self.throttler = AsyncThrottler(rate_limit=self.rate_limit, period=self.period)
|
self.throttler = AsyncThrottler(
|
||||||
|
rate_limit=self.rate_limit, period=self.period
|
||||||
|
)
|
||||||
self._initialized = True
|
self._initialized = True
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
async def get_instance(cls, rate_limit: int = 500, period: int = 60) -> 'GlobalThrottler':
|
async def get_instance(
|
||||||
|
cls, rate_limit: int = 500, period: int = 60
|
||||||
|
) -> "GlobalThrottler":
|
||||||
async with cls._lock:
|
async with cls._lock:
|
||||||
if not cls._instance:
|
if not cls._instance:
|
||||||
cls._instance = cls(rate_limit, period)
|
cls._instance = cls(rate_limit, period)
|
||||||
return cls._instance
|
return cls._instance
|
||||||
|
|
||||||
async def acquire(self):
|
async def acquire(self):
|
||||||
await self.throttler.acquire()
|
await self.throttler.acquire()
|
||||||
|
|
||||||
def release(self):
|
|
||||||
self.throttler.release()
|
|
||||||
|
|
||||||
async def __aenter__(self):
|
async def __aenter__(self):
|
||||||
await self.acquire()
|
await self.acquire()
|
||||||
return self
|
return self
|
||||||
|
|
||||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||||
self.release()
|
# asyncio_throttle.Throttler doesn't have a release() method
|
||||||
|
# It automatically manages rate limiting via acquire()
|
||||||
|
pass
|
||||||
Reference in new issue
Block a user