python async函数_Python-Asynci中的异步函数调节

import asyncio

import contextlib

import collections

import time

from types import TracebackType

from typing import Dict, Optional, Type

try: # Python 3.7

base = contextlib.AbstractAsyncContextManager

_current_task = asyncio.current_task

except AttributeError:

base = object # type: ignore

_current_task = asyncio.Task.current_task # type: ignore

class AsyncLeakyBucket(base):

"""A leaky bucket rate limiter.

Allows up to max_rate / time_period acquisitions before blocking.

time_period is measured in seconds; the default is 60.

"""

def __init__(

self,

max_rate: float,

time_period: float = 60,

loop: Optional[asyncio.AbstractEventLoop] = None

) -> None:

self._loop = loop

self._max_level = max_rate

self._rate_per_sec = max_rate / time_period

self._level = 0.0

self._last_check = 0.0

# queue of waiting futures to signal capacity to

self._waiters: Dict[asyncio.Task, asyncio.Future] = collections.OrderedDict()

def _leak(self) -> None:

"""Drip out capacity from the bucket."""

if self._level:

# drip out enough level for the elapsed time since

# we last checked

elapsed = time.time() - self._last_check

decrement = elapsed * self._rate_per_sec

self._level = max(self._level - decrement, 0)

self._last_check = time.time()

def has_capacity(self, amount: float = 1) -> bool:

"""Check if there is enough space remaining in the bucket"""

self._leak()

requested = self._level + amount

# if there are tasks waiting for capacity, signal to the first

# there there may be some now (they won't wake up until this task

# yields with an await)

if requested < self._max_level:

for fut in self._waiters.values():

if not fut.done():

fut.set_result(True)

break

return self._level + amount <= self._max_level

async def acquire(self, amount: float = 1) -> None:

"""Acquire space in the bucket.

If the bucket is full, block until there is space.

"""

if amount > self._max_level:

raise ValueError("Can't acquire more than the bucket capacity")

loop = self._loop or asyncio.get_event_loop()

task = _current_task(loop)

assert task is not None

while not self.has_capacity(amount):

# wait for the next drip to have left the bucket

# add a future to the _waiters map to be notified

# 'early' if capacity has come up

fut = loop.create_future()

self._waiters[task] = fut

try:

await asyncio.wait_for(

asyncio.shield(fut),

1 / self._rate_per_sec * amount,

loop=loop

)

except asyncio.TimeoutError:

pass

fut.cancel()

self._waiters.pop(task, None)

self._level += amount

return None

async def __aenter__(self) -> None:

await self.acquire()

return None

async def __aexit__(

self,

exc_type: Optional[Type[BaseException]],

exc: Optional[BaseException],

tb: Optional[TracebackType]

) -> None:

return None

请注意，我们偶然地从bucket中泄漏容量，不需要运行单独的异步任务来降低级别；相反，在测试是否有足够的剩余容量时，容量会泄漏出去。在

请注意，等待容量的任务保存在有序字典中，当可能有容量再次空闲时，第一个仍在等待的任务将提前唤醒。在

您可以将其用作上下文管理器；尝试在bucket满块时获取它，直到再次释放足够的容量：

^{pr2}$

也可以直接调用acquire()：await bucket.acquire() # blocks until there is space in the bucket

或者您可以先测试是否有空间：if bucket.has_capacity():

# reject a request due to rate limiting

请注意，您可以通过增加或减少“滴入”桶中的量，将某些请求计数为“较重”或“较轻”：await bucket.acquire(10)

if bucket.has_capacity(0.5):

但是，一定要小心；当混合大滴和小滴时，当达到或接近最大速率时，小滴往往在大滴之前流动，因为有更大的可能性，即在有足够的自由容量来容纳较大的滴液之前有足够的自由容量。在

演示：>>> import asyncio, time

>>> bucket = AsyncLeakyBucket(5, 10)

>>> async def task(id):

... await asyncio.sleep(id * 0.01)

... async with bucket:

... print(f'{id:>2d}: Drip! {time.time() - ref:>5.2f}')

...

>>> ref = time.time()

>>> tasks = [task(i) for i in range(15)]

>>> result = asyncio.run(asyncio.wait(tasks))

0: Drip! 0.00

1: Drip! 0.02

2: Drip! 0.02

3: Drip! 0.03

4: Drip! 0.04

5: Drip! 2.05

6: Drip! 4.06

7: Drip! 6.06

8: Drip! 8.06

9: Drip! 10.07

10: Drip! 12.07

11: Drip! 14.08

12: Drip! 16.08

13: Drip! 18.08

14: Drip! 20.09

bucket在一开始就被迅速地填满，从而使其余的任务分布得更均匀；每2秒钟就释放出足够的容量来处理另一个任务。在

最大突发大小等于最大速率值，在上面的演示中设置为5。如果不希望允许爆发，请将“最大速率”设置为1，将“时间段”设置为“滴流之间的最短时间”：>>> bucket = AsyncLeakyBucket(1, 1.5) # no bursts, drip every 1.5 seconds

>>> async def task():

... async with bucket:

... print(f'Drip! {time.time() - ref:>5.2f}')

...

>>> ref = time.time()

>>> tasks = [task() for _ in range(5)]

>>> result = asyncio.run(asyncio.wait(tasks))

Drip! 0.00

Drip! 1.50

Drip! 3.01

Drip! 4.51

Drip! 6.02