Question

我目前正在实现TCP套接字协议。该协议要求每五分钟发送一次心跳消息。我正在Python中使用asyncio实现协议。下面的源代码是一个程序，该程序连接到localhost：8889，发送问候，并在1秒后断开套接字。在这种情况下，连接将在一秒钟后断开（如果确实发生这种情况，则说明网络已关闭或服务器已断开连接）。问题是send_heartbeat函数等待5分钟，而不知道套接字已关闭。我想立即取消协程，而不是等套接字断开连接后再等待5分钟。最好的方法是什么？

import asyncio


async def run(host: str, port: int):
    while True:
        try:
            reader, writer = await asyncio.open_connection(host, port)

        except OSError as e:
            print('connection failed:', e)
            await asyncio.sleep(0.5)
            continue

        await asyncio.wait([
            handle_stream(reader, writer),
            send_heartbeat(reader, writer),
        ], return_when=asyncio.FIRST_COMPLETED)  # will stop after 1 second

        writer.close()  # close socket after 1 second
        await writer.wait_closed()


async def handle_stream(reader, writer):
    writer.write(b'hello\n')  # will success because socket is alive
    await writer.drain()

    await asyncio.sleep(1)


async def send_heartbeat(reader, writer):
    while True:
        await asyncio.sleep(300)

        heartbeat_message = b'heartbeat\n'
        writer.write(heartbeat_message)  # will fail because socket is already closed after 1 second
        await writer.drain()


if __name__ == '__main__':
    asyncio.run(run('127.0.0.1', 8889))

Answer 1

您可以通过取消执行睡眠的任务来取消睡眠。将send_heartbeat创建为单独的任务可确保在您等待handle_stream的同时并行运行：

async def run(host: str, port: int):
    while True:
        ...
        heartbeat = asyncio.create_task(send_heartbeat(reader, writer))
        try:
            await handle_stream(reader, writer)
        finally:
            heartbeat.cancel()
            writer.close()
        await writer.wait_closed()

顺便说一句，由于您正在writer.drain()中等待handle_stream，因此无法保证handle_stream将始终在1秒内完成。在这里，您可能想避免流失，或者可以在等待handle_stream(...)时使用asyncio.wait_for。

asyncio tcp套接字：关闭套接字后如何取消asyncio.sleep（）？

1 个答案: