本文详细阐述了如何在FastAPI的同一个异步事件循环中，通过正确利用其`lifespan`上下文管理器，同时启动并管理多个异步TCP服务器。核心在于理解`yield`在`lifespan`中的作用，将TCP服务器作为后台任务在应用启动阶段（`yield`之前）调度，并实现优雅的停机机制。通过代码示例，展示了FastAPI、异步TCP服务器与WebSocket的协同工作，实现了数据从TCP到WebSocket的转发。
在FastAPI中运行异步TCP服务器：核心原理与实践
在现代异步应用开发中，常常需要将不同的服务类型（如HTTP API和自定义协议的TCP服务器）集成到同一个应用程序中。FastAPI以其高性能和异步特性，为构建此类复合应用提供了强大的基础。本文将深入探讨如何在FastAPI应用中，利用其lifespan上下文管理器，同时启动并管理多个异步TCP服务器，实现数据从TCP到WebSocket的无缝转发。
FastAPI Lifespan机制解析
FastAPI提供了lifespan上下文管理器，用于在应用启动和关闭时执行特定的初始化和清理任务。其工作原理基于Python的异步上下文管理器协议，通过yield关键字将应用生命周期划分为两个主要阶段：
- 启动阶段 (Startup Phase)：yield 之前的代码块会在FastAPI应用完全启动并准备好接受请求之前执行。这是进行资源初始化、数据库连接、启动后台任务等操作的理想位置。
- 关闭阶段 (Shutdown Phase)：yield 之后的代码块会在FastAPI应用开始关闭时执行。这通常用于清理资源、关闭连接、停止后台任务等操作。
最初尝试将TCP服务器启动逻辑放置在yield之后，导致TCP服务器未能成功启动。这是因为yield之后的代码仅在应用关闭时执行，而非启动时。因此，要使TCP服务器随FastAPI应用一同启动，必须将其启动逻辑放置在yield之前。
正确启动异步TCP服务器
为了在FastAPI应用启动时同时运行异步TCP服务器，我们需要遵循以下步骤：
-
定义异步TCP服务器逻辑：创建一个异步函数来处理TCP客户端连接，并启动TCP服务器。例如，start_tcp_server 函数将监听指定端口，并通过handle_client处理每个连接。
# server.py import asyncio import globals # 假设 globals.py 包含 WebSocketManager 实例 async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter): """处理单个TCP客户端连接，接收数据并广播到WebSocket。""" try: while True: data = await reader.read(1024) if not data: break # 将接收到的数据通过WebSocket广播 await globals.websocket_manager.broadcast(data.decode('utf-8')) except Exception as e: print(f"TCP handle_client error: {e}") finally: writer.close() await writer.wait_closed() async def start_tcp_server(port: int): """启动一个异步TCP服务器监听指定端口。""" print(f"Starting TCP server on port {port}...") server = await asyncio.start_server(handle_client, '0.0.0.0', port) async with server: await server.serve_forever() -
在lifespan中调度TCP服务器：在startup_event函数中，使用asyncio.create_task()将每个TCP服务器的启动函数包装成一个独立的异步任务。这些任务会在yield之前被调度，从而与FastAPI应用同时启动。
# main.py (部分代码) from fastapi import FastAPI, WebSocket import asyncio from contextlib import asynccontextmanager import globals # 假设 globals.py 包含 WebSocketManager 实例 from server import start_tcp_server # 导入TCP服务器启动函数 @asynccontextmanager async def startup_event(app: FastAPI): print("Starting TCP servers...") ports = [8001, 8002, 8003] # 在 yield 之前启动 TCP 服务器任务 # asyncio.create_task() 确保这些服务器在后台运行，不会阻塞 FastAPI 的启动 servers = [asyncio.create_task(start_tcp_server(port)) for port in ports] yield # FastAPI 应用在此处开始接受请求 # 应用关闭时，可以在这里执行清理工作，例如停止TCP服务器 print("Shutting down TCP servers...") # TODO: 实现优雅停机逻辑，发送停止信号给服务器 for task in servers: task.cancel() # 取消任务 await asyncio.gather(*servers, return_exceptions=True) # 等待任务完成取消 app = FastAPI(lifespan=startup_event) @app.websocket("/ws") async def websocket_endpoint(websocket: WebSocket): """处理WebSocket连接，并将其加入到WebSocket管理器。""" print("About to connect to websocket") await globals.websocket_manager.connect(websocket) print("WebSocket connected:", websocket) try: while True: # 保持连接活跃，或处理来自客户端的WebSocket消息 await websocket.receive_text() except Exception as e: print(f"WebSocket Error: {e}") finally: globals.websocket_manager.disconnect(websocket) # globals.py (辅助文件) import threading from websocket_manager import WebSocketManager data_storage = {} data_lock = threading.Lock() # 注意：在异步环境中，更推荐使用 asyncio.Lock websocket_manager = WebSocketManager() # websocket_manager.py (辅助文件) from fastapi import WebSocket from typing import List class WebSocketManager: """管理活跃的WebSocket连接，并提供广播功能。""" def __init__(self): self.active_connections: List[WebSocket] = [] async def connect(self, websocket: WebSocket): await websocket.accept() self.active_connections.append(websocket) def disconnect(self, websocket: WebSocket): self.active_connections.remove(websocket) async def broadcast(self, data: str): # 遍历所有连接并发送数据，处理可能的断开连接 disconnected_connections = [] for connection in self.active_connections: try: await connection.send_text(data) except Exception: disconnected_connections.append(connection) for connection in disconnected_connections: self.active_connections.remove(connection)
优雅停机与资源管理
仅仅取消任务可能不足以实现优雅的停机。一个更健壮的方法是为TCP服务器引入一个内部停止信号。这样，服务器可以在收到信号后完成当前操作并干净地关闭。
以下是一个实现优雅停机机制的示例：
import asyncio
from contextlib import asynccontextmanager
from fastapi import FastAPI
class Server:
"""一个带有停止信号的通用异步服务器示例。"""
def __init__(self):
self._stop = False
async def start_tcp_server(self, port: int):
"""模拟一个持续运行的TCP服务器，可通过_stop标志停止。"""
print(f"TCP server on port {port} started. (Simulated)")
while not self._stop:
# 模拟服务器工作，例如监听连接、处理数据
# 实际的TCP服务器会在这里调用 asyncio.start_server 并 serve_forever
await asyncio.sleep(1) # 模拟工作间隔
print(f"TCP server on port {port} stopped gracefully. (Simulated)")
def stop(self):
"""设置停止标志，通知服务器停止运行。"""
self._stop = True
@asynccontextmanager
async def startup_event(app: FastAPI):
print("Starting TCP servers...")
ports = [8001, 8002, 8003]
# 创建一个 Server 实例来管理所有TCP服务器的停止信号
server_manager = Server()
# 启动TCP服务器任务
servers = [asyncio.create_task(server_manager.start_tcp_server(port)) for port in ports]
yield # FastAPI 应用在此处开始接受请求
print("Shutting down TCP servers...")
# 在应用关闭时，发送停止信号给所有服务器
server_manager.stop()
# 等待所有服务器任务完成其清理工作
await asyncio.gather(*servers)
print("All TCP servers shut down.")
app = FastAPI(lifespan=startup_event)
# 假设这里有其他 FastAPI 路由和 WebSocket 终结点
# 例如，可以集成上面提到的 websocket_endpoint在这个改进的例子中：
- Server 类包含一个_stop标志，用于控制服务器的运行状态。
- start_tcp_server方法会持续运行直到_stop被设置为True。
- 在lifespan的关闭阶段（yield之后），调用server_manager.stop()来触发所有服务器的优雅停机，然后await asyncio.gather(*servers)等待这些任务完成。
注意事项与最佳实践
- 错误处理：在TCP服务器的handle_client和WebSocket的websocket_endpoint中，务必加入健壮的错误处理机制，例如try...except...finally块，以防止单个连接的错误导致整个服务器崩溃。
- 共享状态管理：如果FastAPI应用和TCP服务器需要共享数据（如globals.py中的data_storage），请确保使用适当的同步原语。在异步环境中，推荐使用asyncio.Lock而非threading.Lock，以避免阻塞事件循环。
- 日志记录：为服务器的启动、连接、数据传输和关闭等关键事件添加详细的日志，这对于调试和监控至关重要。
- 端口管理：确保TCP服务器监听的端口（如8001、8002、8003）与FastAPI主应用监听的端口（通常是8000）不冲突。
- 资源限制：考虑TCP服务器可能处理的大量并发连接。根据系统资源和预期负载，可能需要调整文件描述符限制等操作系统参数。
总结
通过正确理解并利用FastAPI的lifespan上下文管理器，我们可以有效地在同一个异步事件循环中集成和管理多种类型的服务，如HTTP API和自定义TCP服务器。关键在于在应用的启动阶段（yield之前）将异步TCP服务器作为非阻塞的后台任务调度，并在关闭阶段（yield之后）实现优雅的停机逻辑。这种模式不仅提高了资源利用率，也简化了复合应用的部署和管理。
