Python Asyncio 库之从ChatGPT Bug了解Cancel机制

# 对应的redis-py版本为4.5.2（新推出的修复版本4.5.3仍然有该问题,4.5.3只修复了PIPE）
import asyncio
from redis.asyncio import Redis


async def main():
    async with Redis(single_connection_client=True) as r:

        await r.set('foo', 'foo')
        await r.set('bar', 'bar')

        t = asyncio.create_task(r.get('foo'))
        await asyncio.sleep(0)  # 模拟IO
        t.cancel()
        try:
            await t
            print('try again, we did not cancel the task in time')
        except asyncio.CancelledError as e:
            print('managed to cancel the task, connection is left open with unread response')

        print('bar:', await r.get('bar'))
        print('ping:', await r.ping())
        print('foo:', await r.get('foo'))


if __name__ == '__main__':
    asyncio.run(main())

# redis.asyncio.client
class Redis:
    async def execute_command(self, *args, **options):
        """Execute a command and return a parsed response"""
        await self.initialize()
        pool = self.connection_pool
        command_name = args[0]
        conn = self.connection or await pool.get_connection(command_name, **options)

        if self.single_connection_client:
            await self._single_conn_lock.acquire()
        try:
            return await conn.retry.call_with_retry(
                lambda: self._send_command_parse_response(
                    conn, command_name, *args, **options
                ),
                lambda error: self._disconnect_raise(conn, error),
            )
        finally:
            if self.single_connection_client:
                self._single_conn_lock.release()
            if not self.connection:
                await pool.release(conn)

    async def _send_command_parse_response(self, conn, command_name, *args, **options):
        await conn.send_command(*args)
        return await self.parse_response(conn, command_name, **options)

import asyncio


async def main():
    # 针对于worker来说是只写的
    write_queue = asyncio.Queue()
    # 针对于worker来说是只读的
    read_queue = asyncio.Queue()

    async def porter():
        """搬运工负责把write_queue中的数据写入到read_queue中"""
        while True:
            item = await write_queue.get()
            await read_queue.put(item)
            write_queue.task_done()

    asyncio.create_task(porter())

    async def worker(cnt):
        """模拟Redis的从写入队列写入数据，并从只读队列中读取数据"""
        await write_queue.put(cnt)
        print(f"worker:{cnt} get num :{await read_queue.get()}")

    # 模拟取消操作
    t = asyncio.create_task(worker(1))
    await asyncio.sleep(0)
    t.cancel()
    try:
        await t
        print("no cancel")
    except asyncio.CancelledError:
        print("cancel")

    await worker(2)
    await worker(3)


if __name__ == '__main__':
    asyncio.run(main())

import asyncio


async def main():
    # 针对于worker来说是只写的
    write_queue = asyncio.Queue()
    # 针对于worker来说是只读的
    read_queue = asyncio.Queue()

    async def porter():
        """搬运工负责把write_queue中的数据写入到read_queue中"""
        while True:
            item = await write_queue.get()
            await read_queue.put(item)
            write_queue.task_done()

    asyncio.create_task(porter())

    async def worker(cnt):
        """模拟Redis的从写入队列写入数据，并从只读队列中读取数据"""
        await write_queue.put(cnt)
        print(f"worker:{cnt} get num :{await read_queue.get()}")

    # 模拟取消操作
    t = asyncio.create_task(worker(1))
    await asyncio.sleep(0.1)
    t.cancel()
    try:
        await t
        print("no cancel")
    except asyncio.CancelledError:
        print("cancel")

    await worker(2)
    await worker(3)


if __name__ == '__main__':
    asyncio.run(main())

worker:1 get num :1
no cancel
worker:2 get num :2
worker:3 get num :3

import asyncio


async def main():
    t = asyncio.create_task(asyncio.sleep(0))
    print(t.done())    # False
    await asyncio.sleep(0.1)
    print(t.done())    # True


if __name__ == '__main__':
    asyncio.run(main())

class Future:
    ... # 省略其他方法
    def cancel(self, msg=None):
        self.__log_traceback = False
        if self._state != _PENDING:
            return False
        self._state = _CANCELLED
        self._cancel_message = msg
        self.__schedule_callbacks()
        return True

import asyncio


async def main():
    async def demo():
        print("demo函数运行了！")
        await asyncio.sleep(10)
    t = asyncio.create_task(demo())
    t.cancel()
    print(t.cancelled())
    try:
        await t
    except asyncio.CancelledError:
        print("demo已经被取消了")
    print(t.cancelled())

if __name__ == '__main__':
    asyncio.run(main())

async def main():
    async def demo():
        try:
            await asyncio.sleep(0)
        except asyncio.CancelledError as e:
            await asyncio.sleep(2)
            raise e

    t = asyncio.create_task(demo())
    await asyncio.sleep(0)
    t.cancel()
    start_cancel_time = time.time()
    print(t.cancelled())
    await asyncio.sleep(0.5)
    print(t.cancelled())
    try:
        await t
    except asyncio.CancelledError as e:
        print(t.cancelled())
    print(f"取消操作的执行时间为：{time.time() - start_cancel_time}")

if __name__ == '__main__':
    asyncio.run(main())

async def main():
    async def demo():
        try:
            await asyncio.sleep(0)
        except asyncio.CancelledError as e:
            print("任务已经被取消了")
            await asyncio.sleep(2)

    t = asyncio.create_task(demo())
    await asyncio.sleep(0)
    t.cancel()
    print(t.cancelled())
    print("取消的标记为： ", t._must_cancel)
    await asyncio.sleep(0)
    try:
        await t
        print("没有收到取消操作, 协程的取消结果为：", t.cancelled())
    except asyncio.CancelledError as e:
        print("已经收到取消操作, 协程的取消结果为：", t.cancelled())

if __name__ == '__main__':
    asyncio.run(main())

class Task(futures._PyFuture):
    def __init__(self, coro, *, loop=None, name=None, context=None):
        ...
        self._must_cancel = False
        self._fut_waiter = None
        ...
    
    def cancel(self, msg=None):
        ...
        if self.done():
            return False
        if self._fut_waiter is not None:
            if self._fut_waiter.cancel(msg=msg):
                return True
        self._must_cancel = True
        self._cancel_message = msg
        return True

    def __step(self, exc=None):
        if self.done():
            raise exceptions.InvalidStateError(
                f'_step(): already done: {self!r}, {exc!r}')
        if self._must_cancel:
            if not isinstance(exc, exceptions.CancelledError):
                exc = self._make_cancelled_error()
            self._must_cancel = False
        coro = self._coro
        self._fut_waiter = None

        _enter_task(self._loop, self)
        # Call either coro.throw(exc) or coro.send(None).
        try:
            if exc is None:
                # We use the `send` method directly, because coroutines
                # don't have `__iter__` and `__next__` methods.
                result = coro.send(None)
            else:
                result = coro.throw(exc)
        except StopIteration as exc:
            if self._must_cancel:
                # Task is cancelled right before coro stops.
                self._must_cancel = False
                super().cancel(msg=self._cancel_message)
            else:
                super().set_result(exc.value)
        except exceptions.CancelledError as exc:
            # Save the original exception so we can chain it later.
            self._cancelled_exc = exc
            super().cancel()  # I.e., Future.cancel(self).
        except (KeyboardInterrupt, SystemExit) as exc:
            super().set_exception(exc)
            raise
        except BaseException as exc:
            super().set_exception(exc)
        else:
            # 这里没有取消操作的相关逻辑，这部分不属于本文的主体，直接省略，知道会通过
            # `self._loop.call_sonn`调用`self.__step`使`Asyncio.Task`能继续往下走
            # 如果如果发现有其他的`coro`或者`Future`对象在等待时就把他们赋值到`self._fut_waiter`就可以了
            # 具体内容请参考《Python的可等待对象在Asyncio的作用》文章
                ...
                result._asyncio_future_blocking = False
                result.add_done_callback(self.__wakeup, context=self._context)
                self._fut_waiter = result
                if self._must_cancel:
                    if self._fut_waiter.cancel(msg=self._cancel_message):
                        self._must_cancel = False
                ...
        finally:
            _leave_task(self._loop, self)
            self = None  # Needed to break cycles when an exception occurs.
            
    def __wakeup(self, future):
        try:
            future.result()
        except BaseException as exc:
            # This may also be a cancellation.
            self.__step(exc)
        else:
            self.__step()
        self = None  # Needed to break cycles when an exception occurs.

import asyncio


async def demo():
    await asyncio.Future()  # <--3:发现coro.send返回的对象是asyncio.Future，此时把asyncio.Future赋值到self._fut_waiter中,并等待`asyncio.Future`执行完毕后来调用自己`self.__wakeup方法`从而再回到`self.__step`方法。
    await asyncio.Future()  # <--4:此时被事件循环唤醒，但是又发现是asyncio.Future，与第三步处理类似
    # <-- 后续没有任务语句，直接抛出StopIteration异常，交由`Asyncio.Task`对象处理

async def main():
    t = asyncio.create_task(demo())  # <-- 1:此时创建一个`asyncio.Task`对象
    await t   # <-- 2:第一次调用了`_step`方法


asyncio.run(main())