在FastAPI中异步管理和监控外部服务的启动与关闭

# 示例：存在问题的等待方式
# while not self.protocal.is_startup:
#     pass

import asyncio
import re
from logging import getLogger

logger = getLogger(__name__)

class MyProtocol(asyncio.SubprocessProtocol):
    def __init__(self, started_future: asyncio.Future, exited_future: asyncio.Future):
        # started_future 和 exited_future 用于向外部传递启动和退出信号
        self.started_future = started_future
        self.exited_future = exited_future
        # 定义一个正则表达式来匹配服务启动成功的日志信息
        self.startup_str = re.compile("Server - Started")

    def pipe_data_received(self, fd, data):
        """
        当子进程的管道（stdout/stderr）接收到数据时调用。
        """
        log_data = data.decode().strip() # 解码并清理日志数据
        logger.info(f"Subprocess Output: {log_data}")
        super().pipe_data_received(fd, data)

        # 检查是否已启动，避免重复设置Future
        if not self.started_future.done():
            if re.search(self.startup_str, log_data):
                logger.info("External service startup signal detected!")
                self.started_future.set_result(True) # 设置Future结果，通知服务已启动

    def pipe_connection_lost(self, fd, exc):
        """
        当子进程的管道连接丢失时调用。
        """
        if exc is None:
            logger.debug(f"Pipe {fd} Closed normally.")
        else:
            logger.error(f"Pipe {fd} Closed with error: {exc}")
        super().pipe_connection_lost(fd, exc)

    def process_exited(self):
        """
        当子进程退出时调用。
        """
        logger.info("External service process exited.")
        super().process_exited()
        # 设置exited_future结果，通知服务已退出
        if not self.exited_future.done():
            self.exited_future.set_result(True)

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

logger = getLogger(__name__)

# 定义全局变量以在lifespan函数外部访问transport和protocol
transport: asyncio.SubprocessTransport
protocol: "MyProtocol"

# MyProtocol 类定义（同上文所示）
class MyProtocol(asyncio.SubprocessProtocol):
    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("Server - Started") # 示例：匹配Java服务启动日志

    def pipe_data_received(self, fd, data):
        log_data = data.decode().strip()
        logger.info(f"Subprocess Output (fd={fd}): {log_data}")
        super().pipe_data_received(fd, data)

        if not self.started_future.done(): # 避免重复设置
            if re.search(self.startup_str, log_data):
                logger.info("External service startup signal detected!")
                self.started_future.set_result(True)

    def pipe_connection_lost(self, fd, exc):
        if exc is None:
            logger.debug(f"Pipe {fd} Closed normally.")
        else:
            logger.error(f"Pipe {fd} Closed with error: {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):
    """
    FastAPI应用的生命周期事件管理器。
    在应用启动时执行yield之前的代码，在应用关闭时执行yield之后的代码。
    """
    global transport, protocol
    loop = asyncio.get_running_loop()

    # 创建Future对象，用于等待外部服务的启动和退出信号
    started_future = asyncio.Future(loop=loop)
    exited_future = asyncio.Future(loop=loop)

    # 启动外部子进程，并传入MyProtocol实例
    # 注意：这里使用lambda函数来延迟MyProtocol的实例化，
    # 确保在subprocess_shell调用时才创建protocol实例并传入Future
    transport, protocol = await loop.subprocess_shell(
        lambda: MyProtocol(started_future, exited_future),
        "/start_java_server.sh" # 替换为你的Java服务启动脚本
    )
    logger.info("External service process started.")

    try:
        # 等待外部服务启动成功，设置超时时间防止无限等待
        await asyncio.wait_for(started_future, timeout=15.0) # 增加超时时间以适应实际情况
        logger.info("External service reported startup success.")
    except asyncio.TimeoutError:
        logger.error("External service startup timed out!")
        # 在超时情况下可以考虑杀死子进程或抛出异常
        transport.close()
        raise RuntimeError("External service failed to start in time.")

    # ------ yield 关键字之前的代码在应用启动时执行 ------
    yield # FastAPI应用在此处开始处理请求

    # ------ yield 关键字之后的代码在应用关闭时执行 ------
    logger.info("FastAPI application shutting down, waiting for external service exit.")
    try:
        # 等待外部服务优雅退出，同样设置超时
        await asyncio.wait_for(exited_future, timeout=10.0)
        logger.info("External service reported graceful shutdown.")
    except asyncio.TimeoutError:
        logger.warning("External service did not exit gracefully within timeout. Forcing close.")
    finally:
        # 无论外部服务是否优雅退出，都关闭transport以清理资源
        if transport.is_closing():
            logger.debug("Subprocess transport is already closing.")
        else:
            transport.close()
            logger.info("Subprocess transport closed.")


app = FastAPI(lifespan=lifespan)

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