InboundVerify/server.py

# server.py
#
# 服务模式入口：FastAPI（主线程，uvicorn/asyncio）+ Playwright worker（独立线程）。
# 客户端经 HTTP 触发任务、查状态、下载数据；worker 串行执行任务并跑心跳。
#
# 线程模型（关键）：
#   - 主线程：uvicorn + FastAPI。路由【绝不】访问 Playwright 对象，只经
#     task_queue（投递任务）+ state_store（查状态/任务）+ 文件系统（下载数据）。
#   - worker 线程：runtime.launch_and_prepare（sync_playwright 在此）+ 任务循环，
#     独占所有 page 操作；与主线程仅经 Queue + SQLite 通信。
#   违反"路由不碰 Playwright"会崩（sync 对象跨线程访问）。
#
# 运行：python server.py  （默认监听 0.0.0.0:8000）

import os
import queue
import threading
import time
from contextlib import asynccontextmanager

import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel

from paths import DOWNLOAD_DIR
import state_store
from runtime import (
    HEARTBEAT_INTERVAL,
    TASK_HANDLERS,
    dispatch_task,
    launch_and_prepare,
    run_heartbeat,
)

# 任务队列：元素 (task_id, task_spec)。主线程投递，worker 消费。
task_queue: "queue.Queue" = queue.Queue()
# worker 运行状态（主线程只读，worker 写）
worker_state = {
    "ctx": None,
    "stop": False,
    "thread": None,
    "ready": False,  # launch_and_prepare 完成（各站就绪，可接任务）
    "error": None,  # worker 启动失败原因
}


def _worker_loop():
    """worker 线程：启动 Playwright + 等就绪 + 任务循环（执行任务 + 心跳）。"""
    try:
        ctx = launch_and_prepare()
        worker_state["ctx"] = ctx
        worker_state["ready"] = True
        print(">> [worker] 各站就绪，开始接收任务 ...")
    except Exception as e:
        worker_state["error"] = str(e)
        print(f"❌ [worker] 启动失败: {e}")
        return

    last_heartbeat = 0.0
    while not worker_state["stop"]:
        try:
            task_id, task_spec = task_queue.get(timeout=1)
        except queue.Empty:
            # 空闲时跑心跳
            if time.monotonic() - last_heartbeat >= HEARTBEAT_INTERVAL:
                run_heartbeat(ctx)
                last_heartbeat = time.monotonic()
            continue

        state_store.update_task(task_id, state_store.TASK_RUNNING)
        print(f">> [worker] 执行任务 #{task_id}: {task_spec}")
        status, error = dispatch_task(ctx, task_spec)
        state_store.update_task(task_id, status, error)
        print(f">> [worker] 任务 #{task_id} 完成: {status} {error or ''}")

    try:
        ctx.stop()
    except Exception:
        pass
    worker_state["ready"] = False
    print(">> [worker] 已退出。")


@asynccontextmanager
async def lifespan(_app):
    """服务启停：起 worker 线程 / 通知 worker 停。"""
    t = threading.Thread(target=_worker_loop, daemon=True)
    worker_state["thread"] = t
    t.start()
    yield
    worker_state["stop"] = True
    t.join(timeout=10)


app = FastAPI(title="InboundVerify 服务端", lifespan=lifespan)


class TaskRequest(BaseModel):
    site: str
    kind: str


@app.post("/tasks")
def create_task(req: TaskRequest):
    """提交任务 {site, kind} → 入队，返回 task_id。"""
    if (req.site, req.kind) not in TASK_HANDLERS:
        raise HTTPException(status_code=400, detail=f"无效任务: {req.site}/{req.kind}")
    task_id = state_store.create_task(req.site, req.kind)
    task_queue.put((task_id, {"site": req.site, "kind": req.kind}))
    return {"task_id": task_id}


@app.get("/tasks/{task_id}")
def get_task(task_id: int):
    t = state_store.get_task(task_id)
    if not t:
        raise HTTPException(status_code=404, detail="任务不存在")
    return t


@app.get("/tasks")
def list_tasks(limit: int = 20):
    return state_store.list_tasks(limit)


@app.get("/status")
def get_status():
    """各站登录态 + 数据态（前端状态盘用），另含 worker 就绪状态。"""
    return {
        "worker_ready": worker_state["ready"],
        "worker_error": worker_state["error"],
        "sites": state_store.get_all_status(),
    }


@app.get("/data/{filename}")
def download_data(filename: str):
    """下载 downloads/ 下的数据文件（防路径穿越）。"""
    if not filename or "/" in filename or "\\" in filename or ".." in filename:
        raise HTTPException(status_code=400, detail="非法文件名")
    path = os.path.join(DOWNLOAD_DIR, filename)
    # 双重校验：解析后绝对路径仍在 DOWNLOAD_DIR 内
    if not os.path.abspath(path).startswith(os.path.abspath(DOWNLOAD_DIR) + os.sep):
        raise HTTPException(status_code=400, detail="非法路径")
    if not os.path.isfile(path):
        raise HTTPException(status_code=404, detail="文件不存在")
    return FileResponse(path, filename=filename)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)