src/observability/server.py

import asyncio
import json
import os
import socket
from contextlib import asynccontextmanager, nullcontext
from pathlib import Path

import uvicorn
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.trustedhost import TrustedHostMiddleware
from starlette.responses import FileResponse, Response, StreamingResponse
from starlette.routing import Route

from observability.events import bus

PAGE = Path(__file__).with_name("web") / "index.html"
# 项目的 logo 和小懒的头像，页面直接用这里的文件。
ASSETS = Path(__file__).resolve().parents[2] / "assets"


class EmbeddedServer(uvicorn.Server):
    def capture_signals(self):
        return nullcontext()


@asynccontextmanager
async def serve(port, events=bus, routes=()):
    listener = socket.socket()
    closing = asyncio.Event()

    async def stream(seq):
        yield f"retry: 2000\nevent: run\ndata: {json.dumps({'run_id': events.run_id})}\n\n"
        live = None
        while True:
            changed = events.changed
            if live != events.live_version:
                live = events.live_version
                yield f"event: live\ndata: {json.dumps(events.live, ensure_ascii=False)}\n\n"
            for event in events.after(seq):
                seq = event.seq
                yield f"id: {events.run_id}:{seq}\ndata: {event.data}\n\n"
            if closing.is_set():
                return
            try:
                async with asyncio.timeout(15):
                    await changed.wait()
            except TimeoutError:
                yield ": keepalive\n\n"

    async def records(request):
        # 页面重连时带上最后收到的事件 ID，同一次运行内只补发之后的事件。
        cursor = request.headers.get("last-event-id") or request.query_params.get("after", "")
        run_id, _, seq = cursor.partition(":")
        return StreamingResponse(
            stream(int(seq) if run_id == events.run_id and seq.isdecimal() else 0),
            media_type="text/event-stream", headers={"Cache-Control": "no-store"},
        )

    async def page(request):
        return FileResponse(PAGE, headers={"Cache-Control": "no-store"})

    async def asset(request):
        path = ASSETS / request.path_params["name"]
        if path.parent != ASSETS or path.suffix != ".svg" or not path.is_file():
            return Response(status_code=404)
        return FileResponse(path)

    try:
        # Windows 上这个选项允许多个进程绑定同一端口，只在 POSIX 系统上用于重启后立即复用端口。
        if os.name == "posix":
            listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            # 其他服务监听全部网卡时，本机地址仍然可以绑定成功并截走发给它的本机连接，
            # 所以先确认这个端口上没有服务在接受连接。
            with socket.socket() as probe:
                probe.settimeout(0.5)
                if port and probe.connect_ex(("127.0.0.1", port)) == 0:
                    raise OSError("端口上已有其他服务")
            listener.bind(("127.0.0.1", port))
        except OSError as error:
            raise RuntimeError(
                f"可观测性端口 {port} 无法使用，请修改 OBSERVABILITY_PORT：{error}"
            ) from error
        url = f"http://127.0.0.1:{listener.getsockname()[1]}"
        # 事件包含对话内容，只接受以本机地址发起的访问，其他域名解析到本机时拒绝。
        application = Starlette(
            routes=[Route("/", page), Route("/events", records), Route("/assets/{name}", asset), *routes],
            middleware=[Middleware(TrustedHostMiddleware, allowed_hosts=["127.0.0.1", "localhost"])],
        )
        server = EmbeddedServer(uvicorn.Config(
            application, log_level="error", lifespan="off", timeout_graceful_shutdown=2,
        ))
        serving = asyncio.create_task(server.serve(sockets=[listener]))
        try:
            async with asyncio.timeout(10):
                while not server.started:
                    if serving.done():
                        await serving
                        raise RuntimeError("可观测性服务没有启动")
                    await asyncio.sleep(0.01)
            yield url
        finally:
            closing.set()
            events.wake()
            server.should_exit = True
            await serving
    finally:
        listener.close()