在多个消费者之间共享一个动态启动的工作者

Question

我正在构建一个使用 asyncio 连接到外部事件流的工作类。它是一个单一的流，但有几个消费者可以启用它。目标是仅在一个或多个消费者需要时保持连接。

我的要求如下：

• worker 实例是在消费者第一次需要时动态创建的。
• 当其他消费者需要它时，他们会重复使用同一个工作实例。
• 当最后一个消费者关闭流时，它会清理其资源。

这听起来很简单。但是，启动顺序引起了我的问题，因为它本身是异步的。因此，假设这个接口：

class Stream:
    async def start(self, *, timeout=DEFAULT_TIMEOUT):
        pass
    async def stop(self):
        pass

我有以下几种情况：

场景 1 - 启动时出现异常

• Consumer 1 请求 worker 启动。
• 工作器启动序列开始
• Consumer 2 请求 worker 启动。
• Worker 启动序列引发异常。
• 两个消费者都应该将异常视为调用 start() 的结果。

场景 2 - 部分异步取消

• Consumer 1 请求 worker 启动。
• 工作器启动序列开始
• Consumer 2 请求 worker 启动。
• 消费者 1 被取消。
• 工作器启动序列完成。
• 消费者 2 应该看到一个成功的开始。

场景 3 - 完成异步取消

• Consumer 1 请求 worker 启动。
• 工作器启动序列开始
• Consumer 2 请求 worker 启动。
• 消费者 1 被取消。
• 消费者 2 被取消。
• 因此必须取消工作器启动序列。

我很难涵盖所有场景，而不会遇到任何竞争条件以及裸露的 Future 或 Event 对象的意大利面条式混乱。

---

这是写start()的尝试。它依赖于_worker() 在完成启动序列时设置一个名为self._worker_ready 的asyncio.Event：

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'

似乎可以工作，但它似乎太复杂了，而且真的很难测试：工人有很多 await 点，如果我要尝试所有可能的取消点，会导致组合爆炸和执行命令。

我需要一个更好的方法。我因此想知道：

• 我的要求合理吗？
• 是否有常见的模式来执行此操作？
• 我的问题是否会引起一些代码异味？

Answer 1

您的要求听起来很合理。我会尝试通过将Event 替换为未来（在本例中为任务）来简化start，使用它来等待启动完成并传播其过程中发生的异常（如果有）。比如：

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

在这段代码中，worker 启动与 worker 协程是分开的，并且也被移到了一个单独的任务中。这个单独的任务可以等待，并且不需要专用的Event，但更重要的是，它允许场景 1 和 2 由相同的代码处理。即使有人取消了第一个消费者，worker启动任务也不会被取消——取消只是意味着等待它的消费者少了。

因此，在消费者取消的情况下，await self._worker_startup_task 对其他消费者来说可以正常工作，而在工作人员启动时发生实际异常的情况下，所有其他服务员都会看到相同的异常，因为任务已经完成。

场景 3 应该会自动运行，因为无论出于何种原因，我们总是会取消消费者无法再观察到的启动。如果消费者因为启动本身失败而离开，那么self._worker_startup_task 将完成（有一个例外）并且它的取消将是空操作。如果是因为所有消费者在等待启动的过程中自己都被取消了，那么self._worker_startup_task.cancel()会按照场景3的要求取消启动序列。

其余代码如下所示（未经测试）：

    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)

Answer 2

通过我之前的测试并整合来自@user4815162342 的建议，我想出了一个可重复使用的解决方案：

st = SharedTask(test())
task1 = asyncio.ensure_future(st.wait())
task2 = asyncio.ensure_future(st.wait(timeout=15))
task3 = asyncio.ensure_future(st.wait())

这是正确的：task2 在 15 秒后自行取消。取消任务对test() 没有影响，除非它们都被取消。在这种情况下，最后一个被取消的任务将手动取消 test() 并等待取消处理完成。

如果传递了协程，则仅在第一个任务开始等待时调度。

最后，在完成后等待共享任务只是立即产生结果（似乎很明显，但初始版本没有）。

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

---

重新提出任务异常取消（在 cmets 中提到）是有意的。它允许这种模式：

async def my_task():
    try:
        await do_stuff()
    except asyncio.CancelledError as exc:
        await flush_some_stuff()     # might raise an exception
        raise exc

客户端可以取消共享任务并处理可能导致的异常，无论my_task 是否包含在SharedTask 中，它的工作方式都是一样的。