在FastAPI中优雅地管理和监控外部服务的启动与关闭

# 错误示例：可能导致阻塞
# while not protocal.is_startup:
#     pass

import asyncio
import re
from logging import getLogger
from fastapi import FastAPI

logger = getLogger(__name__)

app = FastAPI()
transport: asyncio.SubprocessTransport
protocol: "MyProtocol" # 提前声明类型，避免循环引用

class MyProtocol(asyncio.SubprocessProtocol):
    startup_str = re.compile(b"Server - Started") # 注意这里正则匹配的是bytes
    is_startup = False
    is_exited = False

    def pipe_data_received(self, fd, data):
        logger.info(f"Received from fd {fd}: {data.decode().strip()}")
        # super().pipe_data_received(fd, data) # 通常不需要在子类中调用super() for data processing

        if not self.is_startup:
            if re.search(self.startup_str, data): # 直接匹配bytes
                self.is_startup = True
                logger.info("External service reported startup success.")

    def pipe_connection_lost(self, fd, exc):
        if exc is None:
            logger.debug(f"Pipe {fd} Closed")
        else:
            logger.error(f"Pipe {fd} connection lost: {exc}")
        super().pipe_connection_lost(fd, exc)

    def process_exited(self):
        logger.info("External service process exited.")
        super().process_exited()
        self.is_exited = True


@app.on_event("startup")
async def startup_event():
    global transport, protocol
    loop = asyncio.get_running_loop()
    # 假设 /start_java_server.sh 是一个可执行脚本，用于启动Java服务
    transport, protocol = await loop.subprocess_shell(MyProtocol, "/start_java_server.sh")
    logger.info("Attempting to start external service...")

    # 等待外部服务启动完成
    while not protocol.is_startup:
        await asyncio.sleep(0.1) # 允许事件循环处理其他任务
    logger.info("External service has successfully started.")


@app.on_event("shutdown")
async def shutdown_event():
    global transport
    if transport:
        logger.info("Closing external service transport...")
        transport.close()
        # 可以添加一个短暂的等待，确保进程有机会响应关闭信号
        # await asyncio.sleep(0.5)
        logger.info("External service transport closed.")

import asyncio
from contextlib import asynccontextmanager
import re
from logging import getLogger
from fastapi import FastAPI

logger = getLogger(__name__)

# 全局变量用于存储transport和protocol实例
transport: asyncio.SubprocessTransport = None
protocol: "MyProtocol" = None

class MyProtocol(asyncio.SubprocessProtocol):
    # 构造函数现在接收Future对象
    def __init__(self, started_future: asyncio.Future, exited_future: asyncio.Future):
        self.started_future = started_future
        self.exited_future = exited_future
        self.startup_str = re.compile(b"Server - Started") # 匹配bytes

    def pipe_data_received(self, fd, data):
        # 记录原始数据，方便调试
        logger.info(f"[{fd}] {data.decode(errors='ignore').strip()}")
        # super().pipe_data_received(fd, data) # 通常不需要在此处调用

        # 仅当started_future未完成时才尝试匹配启动字符串
        if not self.started_future.done():
            if re.search(self.startup_str, data):
                logger.info("Detected 'Server - Started' in logs.")
                self.started_future.set_result(True) # 标记启动成功

    def pipe_connection_lost(self, fd, exc):
        if exc is None:
            logger.debug(f"Pipe {fd} Closed")
        else:
            logger.error(f"Pipe {fd} connection lost with exception: {exc}")
        super().pipe_connection_lost(fd, exc)

    def process_exited(self):
        logger.info("External service process exited.")
        super().process_exited()
        # 标记进程已退出
        if not self.exited_future.done():
            self.exited_future.set_result(True)

@asynccontextmanager
async def lifespan(app: FastAPI):
    global transport, protocol
    loop = asyncio.get_running_loop()

    # 创建Future对象用于信号通知
    started_future = asyncio.Future(loop=loop)
    exited_future = asyncio.Future(loop=loop)

    # 启动外部进程，并将Future对象传递给协议
    logger.info("Starting external service via subprocess_shell...")
    transport, protocol = await loop.subprocess_shell(
        lambda: MyProtocol(started_future, exited_future),
        "/start_java_server.sh" # 替换为你的Java启动脚本路径
    )

    try:
        # 等待外部服务启动，设置5秒超时
        logger.info("Waiting for external service to report startup (timeout: 5s)...")
        await asyncio.wait_for(started_future, timeout=5.0)
        logger.info("External service successfully started and ready.")
    except asyncio.TimeoutError:
        logger.critical("External service startup timed out!")
        # 在超时情况下，可以选择关闭传输并退出应用，或进行其他错误处理
        if transport:
            transport.close()
        raise RuntimeError("External service failed to start within the given timeout.")

    yield # FastAPI应用开始接收请求

    # 应用关闭时执行的代码
    logger.info("FastAPI application shutting down. Attempting to gracefully close external service...")
    if transport:
        transport.close() # 发送关闭信号给子进程

        try:
            # 等待外部服务进程退出，设置5秒超时
            logger.info("Waiting for external service to exit (timeout: 5s)...")
            await asyncio.wait_for(exited_future, timeout=5.0)
            logger.info("External service gracefully exited.")
        except asyncio.TimeoutError:
            logger.warning("External service did not exit gracefully within the given timeout.")
            # 如果进程未退出，可以考虑发送更强制的终止信号，或记录警告

# 将lifespan上下文管理器注册到FastAPI应用
app = FastAPI(lifespan=lifespan)

# 示例路由
@app.get("/")
async def read_root():
    return {"message": "Hello from FastAPI, external service is running!"}