Allow multiple tasks in an async consumer to run in parallel

[nils-van-zuijlen]

I have a use case for which I'd like to have multiple tasks in a single consumer instance run concurrently.

Here is a sample code that would only work if the consumer can simultaneously run multiple tasks

class Consumer(AsyncWebsocketConsumer):
    async def connect(self):
        self.next_message_id_lock = asyncio.Lock()
        self.next_message_id = 0
        self.waiters = {}

    async def send_call(self, message):
        async with self.next_message_id_lock:
            message_id = str(self.next_message_id)
            self.next_message_id += 1
        self.waiters[message_id] = {"event": asyncio.Event(), "answer": None}

        await self.send_json([message_id, message])
        await self.waiters[message_id]["event"].wait()

        res = self.waiters[message_id]["answer"]
        del self.waiters[message_id]
        return res

    async def receive_json(self, content):
        message_id = content[0]
        message = content[1]

        self.waiters[message_id]["answer"] = message
        self.waiters[message_id]["event"].set()

    async def task_foo(self, message):
        res = await self.send_call({"do": "this"})
        print(res)  # do something

The idea of this sample is that a message in this consumer's channel would cause task_foo to run, using send_call to run a procedure on the connected client, and that send_call only returns once an answer was received.

However, due to the way tasks are scheduled in https://github.com/django/channels/blob/main/channels/utils.py#L51, it runs into a deadlock as send_call waits for receive_json to be called, but that will never happen because a task is currently running in that consumer.

I tried replacing the await_many_dispatch function by the following code, and my deadlock went away. Would it possible to enable multi-tasking in async consumers?

async def await_many_dispatch(consumer_callables, dispatch):
    """
    Given a set of consumer callables, awaits on them all and passes results
    from them to the dispatch awaitable as they come in.
    """
    # Start them all off as tasks
    loop = asyncio.get_event_loop()
    tasks = [
        loop.create_task(consumer_callable())
        for consumer_callable in consumer_callables
    ]
    dispatched = []
    try:
        while True:
            # Wait for any of them to complete
            await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            # Find the completed one(s), yield results, and replace them
            for i, task in enumerate(tasks):
                if task.done():
                    result = task.result()
                    dispatched.append(loop.create_task(dispatch(result)))
                    tasks[i] = asyncio.ensure_future(consumer_callables[i]())
    finally:
        # Make sure we clean up tasks on exit
        for task in (*tasks, *dispatched):
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass