302 lines
11 KiB
Python
302 lines
11 KiB
Python
# 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
|
||
from typing import Dict, Optional
|
||
|
||
import uvicorn
|
||
from apscheduler.schedulers.background import BackgroundScheduler
|
||
from apscheduler.triggers.cron import CronTrigger
|
||
from fastapi import FastAPI, HTTPException
|
||
from fastapi.responses import FileResponse
|
||
from pydantic import BaseModel
|
||
|
||
from inbound_verify.paths import DOWNLOAD_DIR, OUTPUT_DIR
|
||
from inbound_verify import state_store
|
||
from inbound_verify.runtime import (
|
||
HEARTBEAT_INTERVAL,
|
||
TASK_HANDLERS,
|
||
dispatch_task,
|
||
launch_and_prepare,
|
||
run_heartbeat,
|
||
)
|
||
|
||
# 全部站点；百世固定下载当天，不可配置偏移
|
||
ALL_SITES = ["顺心", "百世", "中通", "韵达", "安能"]
|
||
CONFIGURABLE_SITES = {"顺心", "中通", "韵达", "安能"}
|
||
|
||
# 任务队列：元素 (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 启动失败原因
|
||
}
|
||
|
||
# 每日定时下载调度器（进程内；job 只往 task_queue 投任务，不碰 Playwright）
|
||
scheduler = BackgroundScheduler(daemon=True)
|
||
|
||
|
||
def _worker_loop():
|
||
"""worker 线程：启动 Playwright + 等就绪 + 任务循环（执行任务 + 心跳）。"""
|
||
try:
|
||
ctx = launch_and_prepare()
|
||
worker_state["ctx"] = ctx
|
||
worker_state["ready"] = True
|
||
# 【P1-2 重启自愈】worker 就绪后清理上轮遗留的 pending/running 僵尸任务
|
||
cleaned = state_store.fail_stale_tasks()
|
||
if cleaned:
|
||
print(
|
||
f">> [worker] 自愈：清理 {cleaned} 条遗留任务（pending/running → failed）"
|
||
)
|
||
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] 已退出。")
|
||
|
||
|
||
def _enqueue_undelivered(site):
|
||
"""定时 job：把该站 undelivered 任务投到队列（worker 串行处理；本线程不碰 Playwright）。"""
|
||
try:
|
||
tid = state_store.create_task(site, "undelivered")
|
||
task_queue.put((tid, {"site": site, "kind": "undelivered"}))
|
||
print(f">> [定时] 投递 {site}/undelivered 任务 #{tid}")
|
||
except Exception as e:
|
||
print(f">> [定时] 投递 {site} 失败: {e}")
|
||
|
||
|
||
def _reschedule_site(site):
|
||
"""按持久化配置（重新）注册或取消该站的每日定时 job。"""
|
||
job_id = f"site_{site}"
|
||
try:
|
||
scheduler.remove_job(job_id)
|
||
except Exception:
|
||
pass
|
||
cfg = state_store.get_all_config().get(site)
|
||
if not cfg or not cfg.get("schedule_enabled") or not cfg.get("schedule_time"):
|
||
return
|
||
try:
|
||
hh, mm = cfg["schedule_time"].split(":")
|
||
scheduler.add_job(
|
||
_enqueue_undelivered,
|
||
CronTrigger(hour=int(hh), minute=int(mm)),
|
||
args=[site],
|
||
id=job_id,
|
||
replace_existing=True,
|
||
)
|
||
print(f">> [定时] 已注册 {site} 每日 {cfg['schedule_time']} 下载")
|
||
except Exception as e:
|
||
print(f">> [定时] 注册 {site} 失败: {e}")
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(_app):
|
||
"""服务启停：起 worker 线程 / 通知 worker 停。"""
|
||
state_store.init_db() # 先建表/迁移状态库，确保早于 worker 就绪的 /api/status 可用
|
||
for site in ALL_SITES: # 按持久化配置注册各站每日定时 job
|
||
_reschedule_site(site)
|
||
scheduler.start()
|
||
print(">> [定时] 调度器已启动")
|
||
t = threading.Thread(target=_worker_loop, daemon=True)
|
||
worker_state["thread"] = t
|
||
t.start()
|
||
yield
|
||
scheduler.shutdown(wait=False)
|
||
print(">> [定时] 调度器已停止")
|
||
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。"""
|
||
# 【P0】后端未就绪时直接拒绝，避免任务在 worker 启动前入队卡死
|
||
if not worker_state["ready"]:
|
||
raise HTTPException(
|
||
status_code=409, detail="后端尚未就绪，请等待各站点登录完成后再操作"
|
||
)
|
||
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(),
|
||
"ingest": state_store.get_all_ingest_state(),
|
||
}
|
||
|
||
|
||
class ConfigRequest(BaseModel):
|
||
expected_offset: Optional[int] = None
|
||
actual_offset: Optional[int] = None
|
||
schedule_enabled: Optional[bool] = None
|
||
schedule_time: Optional[str] = None
|
||
settings: Optional[Dict[str, str]] = None
|
||
|
||
|
||
def _default_cfg():
|
||
return {
|
||
"expected_offset": 0,
|
||
"actual_offset": 0,
|
||
"schedule_enabled": False,
|
||
"schedule_time": "",
|
||
}
|
||
|
||
|
||
@app.get("/config")
|
||
def get_config():
|
||
"""各站配置：应到/实到偏移 + 每日定时（百世偏移恒 0）。"""
|
||
cfg = state_store.get_all_config()
|
||
return {site: cfg.get(site, _default_cfg()) for site in ALL_SITES}
|
||
|
||
|
||
@app.put("/config/{site}")
|
||
def set_config(site: str, req: ConfigRequest):
|
||
if site not in ALL_SITES:
|
||
raise HTTPException(status_code=400, detail=f"未知站点: {site}")
|
||
cur = state_store.get_all_config().get(site, _default_cfg())
|
||
# 应到/实到偏移（百世锁定当天）
|
||
for kind, val in (
|
||
("expected", req.expected_offset),
|
||
("actual", req.actual_offset),
|
||
):
|
||
if val is not None:
|
||
if site == "百世":
|
||
raise HTTPException(
|
||
status_code=400, detail="百世固定下载当天，不可配置偏移"
|
||
)
|
||
state_store.set_offset(site, kind, val)
|
||
# 每日定时（所有站可配，含百世）
|
||
if req.schedule_enabled is not None or req.schedule_time is not None:
|
||
enabled = (
|
||
req.schedule_enabled
|
||
if req.schedule_enabled is not None
|
||
else cur["schedule_enabled"]
|
||
)
|
||
time_str = (
|
||
req.schedule_time if req.schedule_time is not None else cur["schedule_time"]
|
||
)
|
||
state_store.set_schedule(site, enabled, time_str)
|
||
_reschedule_site(site)
|
||
# 站点专属配置（密码/账号/路径…）
|
||
if req.settings:
|
||
for k, v in req.settings.items():
|
||
state_store.set_setting(site, k, v)
|
||
cfg = state_store.get_all_config().get(site, _default_cfg())
|
||
cfg["settings"] = state_store.get_site_settings(site)
|
||
return cfg
|
||
|
||
|
||
@app.get("/config/{site}/settings")
|
||
def get_site_settings_api(site: str):
|
||
"""单站专属配置（密码/账号/路径…；不进 5s 轮询）。"""
|
||
if site not in ALL_SITES:
|
||
raise HTTPException(status_code=400, detail=f"未知站点: {site}")
|
||
return state_store.get_site_settings(site)
|
||
|
||
|
||
REPORT_FILE = "应到未到数据.xlsx"
|
||
|
||
|
||
@app.get("/report")
|
||
def download_report():
|
||
"""下载 output/应到未到数据.xlsx（跑比对生成；未生成 404）。"""
|
||
path = os.path.join(OUTPUT_DIR, REPORT_FILE)
|
||
if not os.path.isfile(path):
|
||
raise HTTPException(status_code=404, detail="报告尚未生成，请先跑比对")
|
||
return FileResponse(path, filename=REPORT_FILE)
|
||
|
||
|
||
@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)
|
||
|
||
|
||
def main():
|
||
"""服务模式入口。传字符串导入路径（规范写法；不开 reload/workers 时进程内 import，行为等价）。"""
|
||
uvicorn.run("inbound_verify.cli.server:app", host="0.0.0.0", port=8000)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|