Skip to content

Instantly share code, notes, and snippets.

@Snawoot
Last active July 19, 2019 15:51
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save Snawoot/2e41ef2af73c14d6bf05aa9180d249ad to your computer and use it in GitHub Desktop.
Save Snawoot/2e41ef2af73c14d6bf05aa9180d249ad to your computer and use it in GitHub Desktop.
import asyncio
from collections import deque
class Ratelimit:
def __init__(self, rate, loop=None):
self._loop = asyncio.get_event_loop() if loop is None else loop
self._delay = 1. / rate
self._last_released = 0
self._waiters = deque()
self._release_scheduled = False
def _schedule_dispatch(self):
self._loop.call_at(self._last_released + self._delay, self._dispatch)
self._release_scheduled = True
def _dispatch(self):
self._release_scheduled = False
while self._waiters:
fut = self._waiters.popleft()
if not fut.cancelled():
fut.set_result(None)
self._last_released = self._loop.time()
break
if self._waiters:
self._schedule_dispatch()
async def wait(self):
t = self._loop.time()
if self._last_released + self._delay < t and not self._release_scheduled:
self._last_released = t
else:
fut = self._loop.create_future()
self._waiters.append(fut)
if not self._release_scheduled:
self._schedule_dispatch()
await fut
async def __aenter__(self):
await self.wait()
return self
async def __aexit__(self, exc, exc_type, tb):
pass
@Snawoot
Copy link
Author

Snawoot commented Jul 17, 2019

Usage

async def worker(no, ratelimit, n):
    for _ in range(n):
        await asyncio.sleep(random.random() * 2)

        async with ratelimit:  # this section will be run no more often than once per 2 seconds
            print(time.time(), 'Worker #%d: Bang!' % no)

async def main():
    ratelimit = Ratelimit(.5)  # once per 2 seconds

    tasks = [
        loop.create_task(worker(no, ratelimit, 10))
            for no in range(5)
    ]
    await asyncio.wait(tasks)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment