Gradually create async tasks and wait for all of them to complete

Question

I am trying to make a program to make a lot of web-socket connections to the server I've created:

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())

The problem is that all the clients do their jobs one after another:

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
...

I've tried to use asyncio.wait, but all the clients start together. I want them to be created gradually and connected to the server immediately once each of them is created. At the same time continuing creating new clients.

What approach should I apply to accomplish this?

Answer 1

Using asyncio.wait is a good approach. You can combine it with asyncio.ensure_future and asyncio.sleep to create tasks gradually:

@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)

---

EDIT: I implemented a FutureSet class that should do what you want. This set can be filled with futures and removes them automatically when they're done. It is also possible to wait for all the futures to complete.

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)

Example:

@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()