FastAPI/Starlette graceful shutdown server-sent events
缘起
假如你用 FastAPI/Starlette框架，写了一段 Server-sent Events
@app.get('/api/my_stream')
async def api_stream():
async def gen():
while 1:
yield "data: {}\n\n"
await asyncio.sleep(1.0)
return StreamingResponse(gen, media_type="text/event-stream", headers={
'X-Accel-Buffering': 'no',
'Cache-Control': "no-cache",
'Connection': "keep-alive",
})
FastAPI 事件
然后你重启了服务器，比如 gunicorn+uvicorn，你会发现这个连接不会断开，一直输出结果。直到 worker 进程超过配置里的 timeout 参数（默认60秒），被 master 强行杀死然后重启。
一开始尝试了 FastAPI 官方的 shutdown 事件，再在代码的 while 1 加一个 if 判断。发现不管用
@app.on_event('shutdown')
def on_server_shutdown():
app.state.running = False
发现这段代码执行是执行了，但是是在杀死worker的时候生效的。得更加提前。用ASGI推荐的lifespan和 http.disconnect 太复杂，放弃。
Linux signal
尝试暴力监听 worker signal
import signal
signal.signal(signal.SIGINT, on_server_shutdown) # ctrl+c
signal.signal(signal.SIGUSR2, on_server_shutdown)
signal.signal(signal.SIGTERM, on_server_shutdown)
signal.signal(signal.SIGWINCH, on_server_shutdown)
signal.signal(signal.SIGHUP, on_server_shutdown)
signal.signal(signal.SIGQUIT, on_server_shutdown)
发现也不管用。研究了半天，似乎是 UvicornWorker.init_signals 的时候给signal.SIG_DFL了。无语ing
uvicorn shutdown 流程
看gunicorn日志：
[2023-11-02 09:50:38 +0800] [25918] [INFO] Shutting down
[2023-11-02 09:50:38 +0800] [25918] [INFO] Waiting for application shutdown.
[2023-11-02 09:50:38 +0800] [25918] [INFO] Waiting for connections to close. (CTRL+C to force quit)
....
[2023-11-02 09:51:18 +0800] [25918] [INFO] Application shutdown complete.
[2023-11-02 09:51:18 +0800] [25918] [INFO] Finished server process [25918]
[2023-11-02 09:51:18 +0800] [25918] [INFO] Worker exiting (pid: 25918)
观察到 Waiting for connections to close 之后卡住。
属于 uvicorn 的 Server.shutdown() 方法
这个方法调用链：
Server.install_signal_handlers 里注册 signal Server.handle_exit 里设置 Server.should_exit = True- 每秒一次的
Server.on_tick 就会打断 Server.main_loop() 死循环
- 调用
Server.shutdown
如果能拿到 Server.should_exit 自行判断就好了，但是 ASGI 是容器无感的，只继续研究
断开http连接
上面日志有一个比较关键：
# Request shutdown on all existing connections.
for connection in list(self.server_state.connections):
connection.shutdown()
await asyncio.sleep(0.1)
其http连接关闭方法实现为：
def shutdown(self):
"""
Called by the server to commence a graceful shutdown.
"""
if self.cycle is None or self.cycle.response_complete:
self.transport.close()
else:
self.cycle.keep_alive = False
这里应该走的是 else 分支。所以解决方法就是去代码里拿到 cycle.keep_alive 这个属性。
ASGI 接口
跟了一会儿，发现ASGI 在 starlette 里如下流程：
Route() 初始化 self.app = request_response(endpoint)Route().handle() 的时候会调用 self.app(scope, receive, send)unicorn 的 run_asgi 会 result = await app(self.scope, self.receive, self.send) 。其中 app 就是 starlette 的 Route() 实例
所以解决方案逐渐明朗了
黎明的前夜
通过 FastAPI/Starlette 请求的 .receive 属性的 __self__ 拿到 unicorn 的 cycle 实例，然后定时判断上面 shutdown 赋值的 self.cycle.keep_alive = False
@app.get('/api/my_stream')
async def api_stream(req: Request):
async def gen():
while 1:
yield "data: {}\n\n"
await asyncio.sleep(1.0)
if getattr(req.receive.__self__, 'keep_alive', None) is False:
break
return StreamingResponse(...
这里用了个 getattr() 是保证这个 hack 在非 unicorn 下代码也能正常跑
git commit . -m "haha" 上机联调，发现坑了。。。我特么。。。。这个 cycle.keep_alive 默认就是 False
扭转 keep_alive
这玩意一直为 False 的原因是，unicorn 的 HttpToolsProtocol.on_headers_complete 赋值过程：
self.cycle = RequestResponseCycle(..., keep_alive=http_version != "1.0")
这里 http_version 可以通过 req.scope['http_version'] 得到，打印了一下，你猜怎么着，还真tm是。。。原因就是厂里反代 nginx的proxy_http_version没配置。这里是个常见的坑，非常影响性能，因为每个请求会生成一个新的 http 连接。
但是也等不急SA去改配置了。于是直接在代码 async def gen(): 前面写死：
req.receive.__self__.keep_alive = True 。提交，再试。。。咦，怎么 keep_alive 依然为 False？？？？明明刚刚赋值了？
继续读代码，发现第二个坑，发现 unicorn 的 RequestResponseCycle.send() 方法里，在构造返回的时候，有一句
for name, value in headers:
elif name == b"connection" and value.lower() == b"close":
self.keep_alive = False
好吧。。。那么解决方法就是，把 keep_alive = True 挪到返回内部。先完成返回构造，再强行改值。
最后的代码
@app.get('/api/my_stream')
async def api_stream(req: Request):
# req.receive.__self__.keep_alive = True # doesn't work here
async def gen():
req.receive.__self__.keep_alive = True # works here
while 1:
yield "data: {}\n\n"
await asyncio.sleep(1.0)
if req.receive.__self__.keep_alive is False:
break
return StreamingResponse(gen, media_type="text/event-stream", headers={
'X-Accel-Buffering': 'no',
'Cache-Control': "no-cache",
'Connection': "keep-alive",
})
测试重启 worker，该连接在1秒后断开，worker平稳重启。完美。
苦逼撸业务的一天又开始了。
)){kind=link}