逐步创建异步任务并等待所有任务完成

时间:2015-10-16 13:51:03

标签: python python-3.x python-asyncio aiohttp

我正在尝试创建一个程序,以便与我创建的服务器建立大量的Web套接字连接:

class WebSocketClient():

    @asyncio.coroutine
    def run(self):
        print(self.client_id, 'Connecting')
        ws = yield from aiohttp.ws_connect(self.url)
        print(self.client_id, 'Connected')
        print(self.client_id, 'Sending the message')
        ws.send_str(self.make_new_message())

        while not ws.closed:
            msg = yield from ws.receive()

            if msg.tp == aiohttp.MsgType.text:
                print(self.client_id, 'Received the echo')
                yield from ws.close()
                break

        print(self.client_id, 'Closed')


@asyncio.coroutine
def make_clients():

    for client_id in range(args.clients):
        yield from WebSocketClient(client_id, WS_CHANNEL_URL.format(client_id=client_id)).run()


event_loop.run_until_complete(make_clients())

问题在于所有客户一个接一个地完成工作:

0 Connecting
0 Connected
0 Sending the message
0 Received the echo
0 Closed
1 Connecting
1 Connected
1 Sending the message
1 Received the echo
1 Closed
...

我尝试使用{{3}},但所有客户端都在一起开始。我希望它们逐渐创建并在创建每个服务器后立即连接到服务器。同时继续创建新客户。

我应该采用什么方法来实现这个目标?

1 个答案:

答案 0 :(得分:7)

使用asyncio.wait是一种很好的方法。您可以将其与asyncio.ensure_futureasyncio.sleep结合使用,逐步创建任务:

@asyncio.coroutine
def make_clients(nb_clients, delay):
    futures = []
    for client_id in range(nb_clients):
        url = WS_CHANNEL_URL.format(client_id=client_id)
        coro = WebSocketClient(client_id, url).run()
        futures.append(asyncio.ensure_future(coro))
        yield from asyncio.sleep(delay)
    yield from asyncio.wait(futures)

编辑:我实现了一个FutureSet类,它应该做你想要的。这个集合可以填充期货,并在完成后自动删除它们。也可以等待所有期货完成。

class FutureSet:

    def __init__(self, maxsize, *, loop=None):
        self._set = set()
        self._loop = loop
        self._maxsize = maxsize
        self._waiters = []

    @asyncio.coroutine
    def add(self, item):
        if not asyncio.iscoroutine(item) and \
           not isinstance(item, asyncio.Future):
            raise ValueError('Expecting a coroutine or a Future')
        if item in self._set:
            return
        while len(self._set) >= self._maxsize:
            waiter = asyncio.Future(loop=self._loop)
            self._waiters.append(waiter)
            yield from waiter
        item = asyncio.async(item, loop=self._loop)    
        self._set.add(item)
        item.add_done_callback(self._remove)

    def _remove(self, item):
        if not item.done():
            raise ValueError('Cannot remove a pending Future')
        self._set.remove(item)
        if self._waiters:
            waiter = self._waiters.pop(0)
            waiter.set_result(None)

    @asyncio.coroutine
    def wait(self):
        return asyncio.wait(self._set)

示例:

@asyncio.coroutine
def make_clients(nb_clients, limit=0):
    futures = FutureSet(maxsize=limit)
    for client_id in range(nb_clients):
        url = WS_CHANNEL_URL.format(client_id=client_id)
        client = WebSocketClient(client_id, url)
        yield from futures.add(client.run())
    yield from futures.wait()