基于令牌桶算法的异步信号量

Question

我想实现一个异步令牌桶算法，它管理何时运行任务(共同例程)。简单地说，这将控制在任何给定时间内运行的协同例程的数量。

我尝试过这个漏桶和信号量，但它们都不是真正的令牌桶算法。目标是在有能力的情况下耗尽桶(即运行协同例程)，否则在运行下一个例程之前等待足够长的时间。

因此，我决定要做的是重写信号量类，以控制在任何时间运行的协同例程有多少，如下所示：

TokenSemaphore.py

import datetime
from asyncio.locks import Semaphore
import asyncio
import collections
from math import floor


class TokenSemaphore(Semaphore):
    """A Semaphore implementation.
    A semaphore manages an internal counter which is decremented by each
    acquire() call and incremented by each release() call. The counter
    can never go below zero; when acquire() finds that it is zero, it blocks,
    waiting until some other thread calls release().
    Semaphores also support the context management protocol.
    The optional argument gives the initial value for the internal
    counter; it defaults to 1. If the value given is less than 0,
    ValueError is raised.
    """    
    def __init__(self, capacity=1, rate=1, loop=None):
        if capacity < 0:
            raise ValueError("Semaphore initial value must be >= 0")
        self._capacity = capacity
        self._rate = rate
        self._time_table = collections.deque(maxlen=self._capacity)
        super().__init__(value=capacity, loop=loop) 

    @property
    def capacity(self):
        return self._capacity

    def _wake_up_next(self):
        while self._waiters:
            waiter = self._waiters.popleft()
            if not waiter.done():
                waiter.set_result(None)
                return

    def has_capacity(self):
        if len(self._time_table) < self._capacity:
            self._time_table.append(datetime.datetime.now())
            return True
        tf = datetime.datetime.now()
        delta = (tf - self._time_table[0]).total_seconds()
        if delta < self._rate:
            return False
        else:
            self._time_table.append(tf)
            return True    
    def locked(self):
        """Returns True if semaphore can not be acquired immediately."""
        return self._capacity == 0    

    async def acquire(self):
        """Acquire a semaphore.
        If the internal counter is larger than zero on entry,
        decrement it by one and return True immediately.  If it is
        zero on entry, block, waiting until some other coroutine has
        called release() to make it larger than 0, and then return
        True.
        """
        while not self.has_capacity():
            fut = self._loop.create_future()
            self._waiters.append(fut)
            try:
                await fut
            except:
                # See the similar code in Queue.get.
                fut.cancel()
                if self._capacity > 0 and not fut.cancelled():
                    self._wake_up_next()
                raise
        self._capacity -= 1
        return True    

    async def release(self):
        """Release a semaphore, incrementing the internal counter by one.
        When it was zero on entry and another coroutine is waiting for it to
        become larger than zero again, wake up that coroutine.
        """
        tf = datetime.datetime.now()
        delta = (tf - self._time_table[-1]).total_seconds()
        result = self._rate * floor(delta)
        sleep_time = 1.0/float(self._rate) - result if result < 1.0/float(self._rate) else 0
        await asyncio.sleep(sleep_time)
        tf = datetime.datetime.now()
        delta = (tf - self._time_table[-1]).total_seconds()
        self._capacity += result
        self._wake_up_next()

请注意，我的release()是async def，因为我认为如果桶中没有足够的令牌，就需要在这里睡觉。信号量的发布不是async def。我有预感这里是我搞砸的地方，但我不确定。

为了测试我的实现，我编写了以下代码：

run.py

import asyncio
import aiohttp
import re
import datetime
from TokenSemaphore import TokenSemaphore
SITE = "https://example.com"

async def myWorker(semaphore):
    await semaphore.acquire()
    print("Successfully acquired the semaphore")
    async with aiohttp.ClientSession() as session:
        async with session.get(SITE, verify_ssl=False) as resp:
            print(resp.status, datetime.datetime.now() - ref, semaphore.capacity)
    print("Releasing Semaphore")
    await semaphore.release()   


async def main(loop):
    mySemaphore = TokenSemaphore(capacity=40, rate=2)
    # mySemaphore = asyncio.Semaphore(value=40)
    tasks = [myWorker(mySemaphore) for _ in range(44)]
    await asyncio.wait(tasks)
    print("Main Coroutine") 

ref = datetime.datetime.now()
loop = asyncio.get_event_loop()
loop.run_until_complete(main(loop))
print("All Workers Completed")
loop.close()

问题

因此，看起来TokenSemaphore可以工作，但是如果它们是容量的话，它不会耗尽桶。我的print语句显示了桶的可用容量，并显示它有足够的容量(即运行更多任务的可用性)。我不明白为什么我的令牌信号量不运行更多的共同例程，而它们的能力足以做到这一点。

$ python run.py 
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
200 0:00:00.177742 20
Releasing Semaphore
200 0:00:00.178944 20
Releasing Semaphore
200 0:00:00.184608 20
Releasing Semaphore
200 0:00:01.103417 20
Releasing Semaphore
200 0:00:01.105539 22
Releasing Semaphore
200 0:00:01.106280 22
Releasing Semaphore
200 0:00:01.106929 22
Releasing Semaphore
200 0:00:01.107701 22
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:01.110719 29
Releasing Semaphore
200 0:00:01.111228 29
Releasing Semaphore
200 0:00:01.111801 29
Releasing Semaphore
200 0:00:01.112366 29
Releasing Semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
Successfully acquired the semaphore
200 0:00:01.116581 25
Releasing Semaphore
200 0:00:01.153321 25
Releasing Semaphore
200 0:00:01.155235 25
Releasing Semaphore
200 0:00:01.155791 25
Releasing Semaphore
200 0:00:01.156530 25
Releasing Semaphore
200 0:00:01.157258 25
Releasing Semaphore
200 0:00:01.221712 25
Releasing Semaphore
200 0:00:01.223267 25
Releasing Semaphore
200 0:00:01.223724 25
Releasing Semaphore
200 0:00:01.224246 25
Releasing Semaphore
200 0:00:01.224745 25
Releasing Semaphore
200 0:00:01.228829 25
Releasing Semaphore
200 0:00:04.326125 25
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:04.361430 30
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:04.910990 29
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:05.440614 28
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:05.974999 27
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:06.516174 26
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:07.051482 25
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:07.601656 24
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:08.147306 23
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:08.682823 22
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:09.216370 21
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:09.752510 20
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:10.302981 19
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:10.843989 18
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:11.384492 17
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:11.939925 16
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:12.485116 15
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:13.016098 14
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:13.554884 13
Releasing Semaphore
Successfully acquired the semaphore
200 0:00:14.096828 12
Releasing Semaphore
Main Coroutine
All Workers Completed

Answer 1

三个问题：

1. 我们期望_time_table达到_capacity的大小，但是在acquire中它会减少。最好是将更改移到_time_table中，而不是has_capacity。
2. 在release中，result计算为0，因此在协同线苏醒后容量不会增加。只需将容量增加1。
3. 通常，您可能希望睡在acquire上而不是release上，这样就不会在执行结束时无缘无故地等待。

看看这个，看看它是否有帮助：

class TokenSemaphore(Semaphore):
    """A Semaphore implementation.
    A semaphore manages an internal counter which is decremented by each
    acquire() call and incremented by each release() call. The counter
    can never go below zero; when acquire() finds that it is zero, it blocks,
    waiting until some other thread calls release().
    Semaphores also support the context management protocol.
    The optional argument gives the initial value for the internal
    counter; it defaults to 1. If the value given is less than 0,
    ValueError is raised.
    """
    def __init__(self, capacity=1, rate=1, loop=None):
        if capacity < 0:
            raise ValueError("Semaphore initial value must be >= 0")
        self._capacity = capacity
        # Tracks of coroutines waiting on acquire.
        self._waiting = 0
        self._rate = rate
        self._time_table = collections.deque(maxlen=self._capacity)
        # Time of last token that was issued.
        self._last_token = None
        super().__init__(value=capacity, loop=loop)

    @property
    def capacity(self):
        return max(self._capacity - self._waiting, 0)

    def locked(self):
        """Returns True if semaphore can not be acquired immediately."""
        return self.capacity == 0

    def _get_sleep_time(self):
        now = datetime.datetime.now()
        token_freq = datetime.timedelta(seconds=(1.0/float(self._rate)))
        if self._last_token is None:
            delta = now - self._time_table[-1]
            sleep_time = token_freq - delta
            self._last_token = now + sleep_time
            return sleep_time.total_seconds()
        elif self._last_token < now:
            self._last_token += token_freq
            return 0
        else:
            self._last_token += token_freq
            return (self._last_token - now).total_seconds()

    async def acquire(self):
        """Acquire a semaphore.
        If the internal counter is larger than zero on entry,
        decrement it by one and return True immediately.  If it is
        zero on entry, block, waiting until some other coroutine has
        called release() to make it larger than 0, and then return
        True.
        """
        print(self._capacity)
        if self.locked():
            self._waiting += 1
            fut = self._loop.create_future()
            self._waiters.append(fut)
            sleep_time = self._get_sleep_time()
            # Schedule the execution.
            await asyncio.sleep(sleep_time)
            try:
                # Wait for the corresponding task that's already executing to
                # finish.
                await fut
            except:
                # See the similar code in Queue.get.
                fut.cancel()
                if self._capacity > 0 and not fut.cancelled():
                    self._wake_up_next()
                raise
            finally:
                self._waiting -= 1
        else:
            self._last_token = None
        self._capacity -= 1
        self._time_table.append(datetime.datetime.now())
        return True

    def _wake_up_next(self):
        while self._waiters:
            waiter = self._waiters.popleft()
            if not waiter.done():
                waiter.set_result(None)
                return

    async def release(self):
        """Release a semaphore, incrementing the internal counter by one.
        When it was zero on entry and another coroutine is waiting for it to
        become larger than zero again, wake up that coroutine.
        """
        self._capacity += 1
        self._wake_up_next()