在python asyncio上獲取第一個可用的鎖/信號量

Question

在python 3.4中使用新的asyncio，如何獲取一組鎖/信號量中可用的第一個鎖/信號量？

我做的方法是使用wait(return_when=FIRST_COMPLETED) ，然后在我設法獲取一個后取消所有仍在等待的acquire() 。 但我擔心這可能會導致細微的錯誤/競爭條件，我覺得有一種更優雅的方式。

import asyncio as aio

@aio.coroutine
def run():
    sem1, sem2 = (aio.Semaphore(), aio.Semaphore())
    print('initial:', sem1, sem2)
    a = aio.async(sleep(sem1, 1)) # acquire sem1
    print('just after sleep:', sem1, sem2)
    done, pending = yield from aio.wait([sem1.acquire(), sem2.acquire()], return_when=aio.FIRST_COMPLETED)
    print('done:', done)
    print('pending:', pending)
    for task in pending:
        task.cancel()
    print('after cancel:', sem1, sem2)
    yield from aio.wait([a])
    print('after wait:', sem1, sem2)

@aio.coroutine
def sleep(sem, i):
    with (yield from sem):
        yield from aio.sleep(i)

if __name__ == "__main__":
    aio.get_event_loop().run_until_complete(run())

上面的代碼給出了（內存地址編輯）：

initial: <asyncio.locks.Semaphore object at 0x1 [unlocked,value:1]> <asyncio.locks.Semaphore object at 0x2 [unlocked,value:1]>
just after sleep: <asyncio.locks.Semaphore object at 0x1 [unlocked,value:1]> <asyncio.locks.Semaphore object at 0x2 [unlocked,value:1]>
done: {Task(<acquire>)<result=True>}
pending: {Task(<acquire>)<PENDING>}
after cancel: <asyncio.locks.Semaphore object at 0x1 [locked,waiters:1]> <asyncio.locks.Semaphore object at 0x2 [locked]>
after wait: <asyncio.locks.Semaphore object at 0x1 [unlocked,value:1]> <asyncio.locks.Semaphore object at 0x2 [locked]>

Answer 1

如果我正確理解您的問題，您希望擁有兩個不同的鎖池，一個允許每個代理X個連接，另一個允許Y個全局連接。 實際上可以很容易地使用單個Semaphore對象：

class asyncio.Semaphore（value = 1，*，loop = None）

信號量管理一個內部計數器，該計數器由每個acquire（）調用遞減，並由每個release（）調用遞增。 計數器永遠不會低於零; 當acquire（）發現它為零時，它會阻塞，等待其他線程調用release（）。

因此，不是使用每個使用默認value 1初始化的Semaphore對象列表來實現池，只需將單個Semaphore的value初始化為您希望能夠同時運行的最大任務數。

proxy_sem = Semaphore(value=5) # 5 connections will be able to hold this semaphore concurrently
global_sem = Semaphore(value=15) # 15 connections will be able to hold this semaphore

然后在您的代碼中，只需在獲取全局信號量之前始終獲取代理信號量

with (yield from proxy_sem):
     with (yield from global_sem):

這樣，在等待特定於代理的鎖定時，您將不會持有全局鎖定，這可能會阻止來自另一個代理的連接，如果它可以獲得全局鎖定，則可以自由運行。

編輯：

這是一個完整的示例，演示了一種完成此操作的方法，而根本不需要特定於代理的鎖。 相反，您為每個代理運行一個協程，所有代理都使用相同的隊列。 代理協同程序僅通過跟蹤它們已啟動的活動任務來限制它們運行的​​並發任務的數量，並且僅在它們低於限制時啟動新任務。 當代理協程啟動任務時，該任務負責獲取全局信號量。 這是代碼：

import asyncio
import random

PROXY_CONN_LIMIT = 5
GLOBAL_CONN_LIMIT = 20
PROXIES = ['1.2.3.4', '1.1.1.1', '2.2.2.2', '3.3.3.3', '4.4.4.4']

@asyncio.coroutine
def do_network_stuff(item, proxy_info):
    print("Got {}. Handling it with proxy {}".format(item, proxy_info))
    # Wait a random amount of time to simulate actual work being done.
    yield from asyncio.sleep(random.randint(1,7))

@asyncio.coroutine
def handle_item(item, proxy_info, global_sem):
    with (yield from global_sem):  # Get the global semaphore
       yield from do_network_stuff(item, proxy_info)

@asyncio.coroutine
def proxy_pool(proxy_info, queue, global_sem):
    tasks = []
    def remove_item(task, *args):
        tasks.remove(task)
    while True:  # Loop infinitely. We'll return when we get a sentinel from main()
        while len(tasks) < PROXY_CONN_LIMIT: # Pull from the queue until we hit our proxy limit
            item = yield from queue.get()
            print(len(tasks))
            if item is None:  # Time to shut down
                if tasks:
                    # Make sure all pending tasks are finished first.
                    yield from asyncio.wait(tasks)
                print("Shutting down {}".format(proxy_info))
                return
            # Create a task for the work item, and add it to our list of
            # tasks.
            task = asyncio.async(handle_item(item, proxy_info, global_sem))
            tasks.append(task)
        # We've hit our proxy limit. Now we wait for at least one task
        # to complete, then loop around to pull more from the queue.
        done, pending = yield from asyncio.wait(tasks,
                                                return_when=asyncio.FIRST_COMPLETED) 
        # Remove the completed tasks from the active tasks list.
        for d in done:
            tasks.remove(d)

@asyncio.coroutine
def main():
    global_sem = asyncio.Semaphore(GLOBAL_CONN_LIMIT)
    queue = asyncio.Queue()
    tasks = []
    # Start the proxy pools.
    for proxy in PROXIES:
        tasks.append(asyncio.async(proxy_pool(proxy, queue, global_sem)))

    # Send work to the proxy pools.
    for i in range(50):
        yield from queue.put(i)

    # Tell the proxy pools to shut down.
    for _ in PROXIES:
        yield from queue.put(None)

    # Wait for them to shut down.
    yield from asyncio.wait(tasks)

if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())

樣本輸出：

0
1
2
3
4
0
1
2
3
4
0
1
2
3
4
0
1
2
3
4
0
1
2
3
4
Got 0. Handling it with proxy 1.2.3.4
Got 1. Handling it with proxy 1.2.3.4
Got 2. Handling it with proxy 1.2.3.4
Got 3. Handling it with proxy 1.2.3.4
Got 4. Handling it with proxy 1.2.3.4
Got 5. Handling it with proxy 1.1.1.1
Got 6. Handling it with proxy 1.1.1.1
Got 7. Handling it with proxy 1.1.1.1
Got 8. Handling it with proxy 1.1.1.1
Got 9. Handling it with proxy 1.1.1.1
Got 10. Handling it with proxy 2.2.2.2
Got 11. Handling it with proxy 2.2.2.2
Got 12. Handling it with proxy 2.2.2.2
Got 13. Handling it with proxy 2.2.2.2
Got 14. Handling it with proxy 2.2.2.2
Got 15. Handling it with proxy 3.3.3.3
Got 16. Handling it with proxy 3.3.3.3
Got 17. Handling it with proxy 3.3.3.3
Got 18. Handling it with proxy 3.3.3.3
Got 19. Handling it with proxy 3.3.3.3
Got 20. Handling it with proxy 4.4.4.4
Got 21. Handling it with proxy 4.4.4.4
Got 22. Handling it with proxy 4.4.4.4
Got 23. Handling it with proxy 4.4.4.4
4
4
4
4
Got 24. Handling it with proxy 4.4.4.4
Got 25. Handling it with proxy 1.2.3.4
Got 26. Handling it with proxy 2.2.2.2
Got 27. Handling it with proxy 1.1.1.1
Got 28. Handling it with proxy 3.3.3.3
3
4
4
4
4
4
Got 29. Handling it with proxy 4.4.4.4
Got 30. Handling it with proxy 4.4.4.4
Got 31. Handling it with proxy 2.2.2.2
Got 32. Handling it with proxy 1.1.1.1
4
4
4
Got 33. Handling it with proxy 1.2.3.4
Got 34. Handling it with proxy 3.3.3.3
Got 35. Handling it with proxy 1.1.1.1
Got 36. Handling it with proxy 2.2.2.2
Got 37. Handling it with proxy 3.3.3.3
3
4
4
4
4
Got 38. Handling it with proxy 1.2.3.4
4
Got 39. Handling it with proxy 1.2.3.4
Got 40. Handling it with proxy 2.2.2.2
Got 41. Handling it with proxy 1.1.1.1
Got 42. Handling it with proxy 3.3.3.3
Got 43. Handling it with proxy 4.4.4.4
2
3
4
4
4
4
Got 44. Handling it with proxy 1.2.3.4
Got 45. Handling it with proxy 1.2.3.4
Got 46. Handling it with proxy 1.2.3.4
Got 47. Handling it with proxy 1.1.1.1
Got 48. Handling it with proxy 4.4.4.4
Got 49. Handling it with proxy 2.2.2.2
3
4
4
4
Shutting down 3.3.3.3
4
Shutting down 2.2.2.2
Shutting down 1.1.1.1
Shutting down 4.4.4.4
Shutting down 1.2.3.4