Python 异步编程与 优雅停机 实践

起因

前阵子写了个基于 asyncio 的消费服务，从队列里拿任务、处理、写库，平时跑得好好的。直到有一天更新部署，习惯性 docker compose restart 下去，服务是重启了，一查数据库发现好几条记录状态停在中间，消息队列里还躺着一堆 unacked 的消息等着重投，等于这批活儿白干了。

最早我还试图自救过一把，在信号处理里直接 loop.stop()，因为当时是用 run_forever 起的循环，结果控制台刷了一屏的 Task was destroyed but it is pending!，比不处理还难看。后来老老实实查了资料才发现，优雅停机这事儿在 asyncio 里其实有一套固定打法，下面是实测可用的写法。

先把信号接住

docker stop 发的是 SIGTERM，Ctrl+C 发的是 SIGINT，Python 默认收到这俩直接就把进程干掉了，任务拦腰斩断，谁也不会跟你打招呼。所以第一步是把信号接住，转成一个 asyncio 事件：

import asyncio
import signal

async def main():
    stop = asyncio.Event()
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, stop.set)

    # ... 后面在循环里等这个 stop

这里用 loop.add_signal_handler 而不是老文章里常见的 signal.signal，因为后者的回调里没法 await，配合 asyncio 很别扭。这玩意儿只能在主线程用，而且 Windows 的 ProactorEventLoop 不支持，本地是 Windows 的话会直接甩你一个 NotImplementedError，没准儿你调试的时候会在这上面懵一会儿，像我一样。

顺带一提，注册了 SIGINT 的 handler 之后，Ctrl+C 就不会再抛 KeyboardInterrupt 了，会被当成普通信号处理，在容器里跑反而省心。

等在跑的任务干完

收到信号只是开始，接下来得让在跑的任务把手头的活儿干完。我的做法是用一个 set 收集所有在跑的 task，退出的时候先等它们自然结束，等不到再砍：

class Worker:
    def __init__(self):
        self.pending: set[asyncio.Task] = set()

    def spawn(self, coro):
        t = asyncio.create_task(coro)
        self.pending.add(t)
        t.add_done_callback(self.pending.discard)

    async def drain(self, timeout: float = 15):  # <--- 按你任务最长的正常耗时调
        if not self.pending:
            return
        print(f"还有 {len(self.pending)} 个任务没收尾，最多等 {timeout} 秒")
        _, stuck = await asyncio.wait(self.pending, timeout=timeout)
        for t in stuck:
            t.cancel()
        if stuck:
            await asyncio.gather(*stuck, return_exceptions=True)

注意这个 timeout 是必须的，不是可选项。不然只要有任务卡死（下游接口挂了、连接断了），进程就永远退不出去，k8s 超过 grace period 直接一记 SIGKILL，前面等的全白等。

另外不是所有任务都值得等。像心跳上报这种，砍了就砍了，单独存个引用，shutdown 时直接 cancel 就行，真正需要等的是写库、调接口这种做一半会留脏数据的。如果任务里有一段“要么不开始、开始了就必须做完”的写入，比如最后落库那一下，可以用 asyncio.shield() 包住它，外面被 cancel 时里面的协程还能继续跑完，当然前提还是进程别被 SIGKILL，那就真没辙了。

再关资源

任务收尾完，最后把 session、连接池这些关干净。关键是顺序：先 drain 再关资源，反了的话任务还在写库，你先把连接池关了，又是一堆报错，别问我怎么知道的。

await worker.drain(timeout=15)
await session.close()      # aiohttp
await engine.dispose()     # sqlalchemy / asyncpg，名字大同小异

拼起来的完整例子

把上面的东西拼一起，一个最小可跑的版本：

import asyncio
import signal
import time

async def handle(item: int):
    print(f"[{time.strftime('%H:%M:%S')}] {item} 开始处理")
    await asyncio.sleep(3)  # 模拟写库、调接口这些耗时操作
    print(f"[{time.strftime('%H:%M:%S')}] {item} 处理完成")

class Worker:
    def __init__(self):
        self.pending: set[asyncio.Task] = set()

    def spawn(self, coro):
        t = asyncio.create_task(coro)
        self.pending.add(t)
        t.add_done_callback(self.pending.discard)

    async def drain(self, timeout: float = 15):
        if not self.pending:
            return
        print(f"还有 {len(self.pending)} 个任务没收尾，最多等 {timeout} 秒")
        _, stuck = await asyncio.wait(self.pending, timeout=timeout)
        for t in stuck:
            t.cancel()
        if stuck:
            await asyncio.gather(*stuck, return_exceptions=True)

async def main():
    worker = Worker()
    stop = asyncio.Event()
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, stop.set)

    n = 0
    while not stop.is_set():
        worker.spawn(handle(n))
        n += 1
        try:
            await asyncio.wait_for(stop.wait(), timeout=1)
        except asyncio.TimeoutError:
            pass

    await worker.drain()
    print("收工")

asyncio.run(main())

直接 python main.py 跑起来，然后 Ctrl+C（或者另开个终端 kill -TERM 一下进程），不出意外的话就能看到它不慌不忙把手头的任务一个一个处理完才退出：

[18:40:02] 9 处理完成
还有 3 个任务没收尾，最多等 15 秒
[18:40:03] 10 处理完成
[18:40:04] 11 处理完成
[18:40:05] 12 处理完成
收工

3.11 之后多了个 asyncio.timeout() 的新写法，更顺一点，但 wait_for 照样能用，老项目不用急着改。

部署那边也得配合

代码写好还不算完，部署平台得给它留出时间。先说 docker，docker stop 和 docker compose restart 默认只等 10 秒，超时直接 SIGKILL，任务要是干不完，compose 里加一行 stop_grace_period: 30s，或者 docker stop -t 30。k8s 对应的是 terminationGracePeriodSeconds，默认 30 秒，另外它的 preStop 和 endpoint 摘除是并行进行的，别指望 preStop 里 sleep 几秒就万事大吉。systemd 的话看 TimeoutStopSec，一个道理。

我一开始以为代码里处理了信号就完事了，结果 docker stop 十秒一到照样强杀，前面的等待全白做。另外，如果你是拿 uvicorn、gunicorn 这类跑的 web 服务，人家框架自带优雅停机，别再自己注册一遍信号处理，两边会打架。

最后

剥掉细节，优雅停机核心就三步：接住信号、等任务干完、关资源。原理一句话就能讲完，但真到生产上，超时兜底、部署平台的 grace period、Windows 的坑，这些边角才是花时间的地方。上面这套抄走改改就能用。