Ndtyjky

I am building a worker class that connects to an external event stream using asyncio. It is a single stream, but several consumer may enable it. The goal is to only maintain the connection while one or more consumer requires it.

My requirements are as follow:

• The worker instance is created dynamically first time a consumer requires it.


• When other consumers then require it, they re-use the same worker instance.


• When the last consumer closes the stream, it cleans up its resources.

This sounds easy enough. However, the startup sequence is causing me issues, because it is itself asynchronous. Thus, assuming this interface:

class Stream:

    async def start(self, *, timeout=DEFAULT_TIMEOUT):

        pass

    async def stop(self):

        pass

I have the following scenarios:

Scenario 1 - exception at startup

• Consumer 1 requests worker to start.


• Worker startup sequence begins


• Consumer 2 requests worker to start.


• Worker startup sequence raises an exception.


• Both consumer should see the exception as the result of their call to start().

Scenario 2 - partial asynchronous cancellation

• Consumer 1 requests worker to start.


• Worker startup sequence begins


• Consumer 2 requests worker to start.


• Consumer 1 gets cancelled.


• Worker startup sequence completes.


• Consumer 2 should see a successful start.

Scenario 3 - complete asynchronous cancellation

• Consumer 1 requests worker to start.


• Worker startup sequence begins


• Consumer 2 requests worker to start.


• Consumer 1 gets cancelled.


• Consumer 2 gets cancelled.


• Worker startup sequence must be cancelled as a result.

I struggle to cover all scenarios without getting any race condition and a spaghetti mess of either bare Future or Event objects.

Here is an attempt at writing start(). It relies on _worker() setting an asyncio.Event named self._worker_ready when it completes the startup sequence:

async def start(self, timeout=None):

    assert not self.closing

    if not self._task:

        self._task = asyncio.ensure_future(self._worker())



    # Wait until worker is ready, has failed, or timeout triggers

    try:

        self._waiting_start += 1

        wait_ready = asyncio.ensure_future(self._worker_ready.wait())



        done, pending = await asyncio.wait(

            [self._task, wait_ready],

            return_when=asyncio.FIRST_COMPLETED, timeout=timeout

        )

    except asyncio.CancelledError:

        wait_ready.cancel()

        if self._waiting_start == 1:

            self.closing = True

            self._task.cancel()

            with suppress(asyncio.CancelledError):

                await self._task    # let worker shutdown

        raise

    finally:

        self._waiting_start -= 1



    # worker failed to start - either throwing or timeout triggering

    if not self._worker_ready.is_set():

        self.closing = True

        self._task.cancel()

        wait_ready.cancel()



        try:

            await self._task        # let worker shutdown

        except asyncio.CancelledError:

            raise FeedTimeoutError('stream failed to start within %ss' % timeout)

        else:

            assert False, 'worker must propagate the exception'

That seems to work, but it seems too complex, and is really hard to test: the worker has many await points, leading to combinatoric explosion if I am to try all possible cancellation points and execution orders.

I need a better way. I am thus wondering:

• Are my requirements reasonable?


• Is there a common pattern to do this?


• Does my question raise some code smell?

Your requirements sound reasonable. I would try to simplify start by replacing Event with a future (in this case a task), using it to both wait for the startup to finish and to propagate exceptions that occur during its course, if any. Something like:

class Stream:

    async def start(self, *, timeout=DEFAULT_TIMEOUT):

        loop = asyncio.get_event_loop()

        if self._worker_startup_task is None:

            self._worker_startup_task = 

                loop.create_task(self._worker_startup())



        self._add_user()

        try:

            await asyncio.shield(asyncio.wait_for(

                self._worker_startup_task, timeout))

        except:

            self._rm_user()

            raise



    async def _worker_startup(self):

        loop = asyncio.get_event_loop()

        await asyncio.sleep(1)      # ...

        self._worker_task = loop.create_task(self._worker())

In this code the worker startup is separated from the worker coroutine, and is also moved to a separate task. This separate task can be awaited and removes the need for a dedicated Event, but more importantly, it allows scenarios 1 and 2 to be handled by the same code. Even if someone cancels the first consumer, the worker startup task will not be canceled - the cancellation just means that there is one less consumer waiting for it.

Thus in case of consumer cancellation, await self._worker_startup_task will work just fine for other consumers, whereas in case of an actual exception in worker startup, all other waiters will see the same exception because the task will have completed.

Scenario 3 should work automatically because we always cancel the startup that can no longer be observed by a consumer, regardless of the reason. If the consumers are gone because the startup itself has failed, then self._worker_startup_task will have completed (with an exception) and its cancellation will be a no-op. If it is because all consumers have been themselves canceled while awaiting the startup, then self._worker_startup_task.cancel() will cancel the startup sequence, as required by scenario 3.

The rest of the code would look like this (untested):

    def __init__(self):

        self._users = 0

        self._worker_startup = None



    def _add_user(self):

        self._users += 1



    def _rm_user(self):

        self._users -= 1

        if self._users:

            return

        self._worker_startup_task.cancel()

        self._worker_startup_task = None

        if self._worker_task is not None:

            self._worker_task.cancel()

            self._worker_task = None



    async def stop(self):

        self._rm_user()



    async def _worker(self):

        # actual worker...

        while True:

            await asyncio.sleep(1)

With my previous tests and integrating suggestions from @user4815162342 I came up with a re-usable solution:

st = SharedTask(test())

task1 = asyncio.ensure_future(st.wait())

task2 = asyncio.ensure_future(st.wait(timeout=15))

task3 = asyncio.ensure_future(st.wait())

This does the right thing: task2 cancels itself after 15s. Cancelling tasks has no effect on test() unless they all get cancelled. In that case, the last task to get cancelled will manually cancel test() and wait for cancellation handling to complete.

If passed a coroutine, it is only scheduled when first task starts waiting.

Lastly, awaiting the shared task after it has completed simply yields its result immediately (seems obvious, but initial version did not).

import asyncio

from contextlib import suppress





class SharedTask:

    __slots__ = ('_clients', '_task')



    def __init__(self, task):

        if not (asyncio.isfuture(task) or asyncio.iscoroutine(task)):

            raise TypeError('task must be either a Future or a coroutine object')

        self._clients = 0

        self._task = task



    @property

    def started(self):

        return asyncio.isfuture(self._task)



    async def wait(self, *, timeout=None):

        self._task = asyncio.ensure_future(self._task)



        self._clients += 1

        try:

            return await asyncio.wait_for(asyncio.shield(self._task), timeout=timeout)

        except:

            self._clients -= 1

            if self._clients == 0 and not self._task.done():

                self._task.cancel()

                with suppress(asyncio.CancelledError):

                    await self._task

            raise



    def cancel(self):

        if asyncio.iscoroutine(self._task):

            self._task.close()

        elif asyncio.isfuture(self._task):

            self._task.cancel()

The re-raising of task exception cancellation (mentioned in comments) is intentional. It allows this pattern:

async def my_task():

    try:

        await do_stuff()

    except asyncio.CancelledError as exc:

        await flush_some_stuff()     # might raise an exception

        raise exc

The clients can cancel the shared task and handle an exception that might arise as a result, it will work the same whether my_task is wrapped in a SharedTask or not.