记一次 Websocket 导致的 Fastapi 内存泄漏

之前用 Fastapi 写了一个通知主动推送的系统，使用 Websocket 来实现实时推送。上线运行了一段时间后发现内存占用以每天三四百兆的速度在上升。想着是小系统，会创建的对象就那几个，就没用 tracemalloc。一开始怀疑是数据库读写的 session 没回收，结果分析了一遍，在创建 session 的地方加上了显式的强制回收也还没解决。一开始根本没往 Websocket 的方面去怀疑，因为在路由里的 finally 是有写回收的。

上代码：

@router.websocket("/")
async def ws(
    websocket: WebSocket,
):
    await websocket.accept()
    logger.info(f"WebSocket connection established")
    try:
        # 添加到连接池
        await add_websocket(websocket)
        # 持续监听消息
        while True:
            try:
                # 接收消息
                data = await websocket.receive_text()
                logger.debug(f"Received message from {staffId}: {data}")
                # 处理消息（只处理心跳等基础消息）
                response = "pong"
                # 发送响应
                if response:
                    await websocket.send_text(response)
            except WebSocketDisconnect:
                logger.info(f"WebSocket disconnected")
                break
            except Exception as e:
                logger.error(f"Error processing message: {str(e)}")
                await websocket.send_text(
                    json.dumps({"status": "error", "message": str(e)})
                )

    except Exception as e:
        logger.error(f"WebSocket error: {str(e)}")
    finally:
        # 从连接池移除
        await remove_websocket(websocket)
        logger.info(f"WebSocket connection closed")

这段代码正常人应该都想不到内存要怎么才会泄漏吧。

但是看 pmap 的数据，大概率就是 Websocket 泄露了。为啥 finally 会没能回收掉连接呢？如果客户端突然关机，没来得及断开连接可能就会导致服务端一直都以为客户端还在线。偏偏我又是客户端 ping 服务端来维持心跳，服务端不会去主动检测客户端的存活，所以会导致内存的泄露。所以解决方案就两个：1. 写一个主动心跳。2. 存储每一个客户端的最后心跳时间，定期删除过期连接。