如何在FastAPI应用中优雅地集成并管理异步TCP服务器

import threading
from websocket_manager import WebSocketManager

# 存储共享数据
data_storage = {}
# 用于数据访问的线程锁
data_lock = threading.Lock()
# WebSocket连接管理器
websocket_manager = WebSocketManager()

from fastapi import WebSocket
from typing import List

class WebSocketManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        """建立WebSocket连接并添加到活动连接列表"""
        await websocket.accept()
        self.active_connections.append(websocket)
        print(f"WebSocket connected: {websocket.client}")

    def disconnect(self, websocket: WebSocket):
        """断开WebSocket连接并从活动连接列表移除"""
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
            print(f"WebSocket disconnected: {websocket.client}")

    async def broadcast(self, data: str):
        """向所有活动WebSocket连接广播数据"""
        for connection in self.active_connections:
            try:
                await connection.send_text(data)
            except Exception as e:
                print(f"Error broadcasting to WebSocket {connection.client}: {e}")
                # 如果发送失败，可以考虑断开该连接
                self.disconnect(connection)

import asyncio
import globals

async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
    """处理单个TCP客户端连接"""
    addr = writer.get_extra_info('peername')
    print(f"TCP client connected from {addr}")
    try:
        while True:
            data = await reader.read(1024) # 读取数据
            if not data:
                break # 客户端断开连接
            decoded_data = data.decode('utf-8', errors='ignore')
            print(f"Received from TCP {addr}: {decoded_data}")
            # 通过WebSocket广播接收到的数据
            await globals.websocket_manager.broadcast(decoded_data)
    except asyncio.CancelledError:
        print(f"TCP client handler for {addr} cancelled.")
    except Exception as e:
        print(f"Error handling TCP client {addr}: {e}")
    finally:
        writer.close()
        await writer.wait_closed()
        print(f"TCP client {addr} disconnected.")

async def create_and_run_tcp_server(port: int):
    """
    创建并运行一个TCP服务器。
    此函数返回一个asyncio.Server实例，
    其serve_forever()方法将作为后台任务运行。
    """
    print(f"Attempting to start TCP server on port {port}...")
    server = await asyncio.start_server(handle_client, '0.0.0.0', port)
    print(f"TCP server started on port {port}")
    # serve_forever()是一个阻塞调用，需要通过create_task在后台运行
    # 并且在需要关闭时，调用server.close()来停止它
    await server.serve_forever()
    return server # 实际上，serve_forever会一直运行，直到被关闭，所以这里通常不会返回

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import asyncio
import globals
from server import create_and_run_tcp_server # 导入TCP服务器创建函数
from contextlib import asynccontextmanager

# 用于存储TCP服务器实例和其运行任务，以便在应用关闭时进行管理
tcp_servers = []
tcp_server_tasks = []

@asynccontextmanager
async def startup_event(app: FastAPI):
    """
    FastAPI应用的生命周期管理器。
    在yield之前启动所有后台服务，在yield之后处理服务关闭。
    """
    print("--- FastAPI Application Startup ---")
    ports = [8001, 8002, 8003] # 定义需要启动的TCP服务器端口

    # 启动TCP服务器
    print(f"Starting TCP servers on ports: {ports}")
    for port in ports:
        # 创建TCP服务器实例
        server_instance = await asyncio.start_server(globals.handle_client, '0.0.0.0', port)
        tcp_servers.append(server_instance)

        # 将服务器的serve_forever方法作为后台任务运行
        task = asyncio.create_task(server_instance.serve_forever())
        tcp_server_tasks.append(task)
        print(f"TCP server task created for port {port}")

    # 应用启动完成，现在可以处理请求
    yield

    # 应用关闭阶段：停止所有TCP服务器
    print("--- FastAPI Application Shutdown ---")
    print("Stopping TCP servers...")
    for server_instance in tcp_servers:
        server_instance.close() # 向TCP服务器发送关闭信号

    # 等待所有TCP服务器任务完成关闭
    # return_exceptions=True 确保即使某个任务关闭失败，其他任务也能继续等待
    await asyncio.gather(*tcp_server_tasks, return_exceptions=True)
    print("All TCP servers stopped gracefully.")
    print("--- FastAPI Application Shutdown Complete ---")


# 创建FastAPI应用实例，并指定lifespan管理器
app = FastAPI(lifespan=startup_event)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """
    FastAPI的WebSocket端点，用于客户端连接。
    """
    print("Attempting to connect to WebSocket...")
    await globals.websocket_manager.connect(websocket)
    print(f"WebSocket connected: {websocket.client}")
    try:
        while True:
            # 保持WebSocket连接活跃，并处理可能接收到的消息
            # 这里我们只是接收，不处理，因为数据流是从TCP到WebSocket
            message = await websocket.receive_text()
            print(f"Received from WebSocket {websocket.client}: {message}")
            # 如果需要，可以将WebSocket接收到的数据转发给TCP服务器
            # await some_tcp_client_writer.write(message.encode())
    except WebSocketDisconnect:
        print(f"WebSocket {websocket.client} disconnected.")
    except Exception as e:
        print(f"WebSocket Error for {websocket.client}: {e}")
    finally:
        globals.websocket_manager.disconnect(websocket)

uvicorn main:app --reload