socks-manager/app.py

"""Flask 应用工厂。"""
import atexit
import os
import logging
import threading
import time
from flask import Flask
from config import Config
from database import db
from models import Instance  # 触发表创建

# ── 全局服务 ───────────────────────────────────────────────────
from engine.instances import InstanceManager
from services.user_service import UserService
from services.backup_service import set_backup_dir

# ── SOCKS5 实例自动 sync 调度 ────────────────────────────────
# 只有 master 进程会跑 (gunicorn -w 1 时 = worker 0)。
# 周期调用 instance_manager.sync_instances(), 它会调 Socks5Server.is_alive()
# 探测死掉的实例并用 DB 配置自动重启, 同时修正 DB.running 字段。
#
# 之前 sync_instances 只能由 API 触发 (/api/instances/sync), 没人调就没人探活,
# "实例莫名停止" 实际是 SOCKS5 线程死掉但没人发现。
#
# 30s 周期足够短, 用户感觉不到抖动; 也不会跟 gunicorn worker timeout (30s)
# 竞争——gunicorn 在 HTTP 请求上下文里 timeout, sync_instances 在主线程跑
# 没有 gunicorn 中间件, 不会触发 worker timeout。
_sync_thread = None
_sync_lock = threading.Lock()


def _start_sync_loop(app):
    """启动后台 sync 线程。线程级单例, 多次调用只启一次。

    gunicorn -w N 时, 每个 worker 都会调 create_app() 一次, 但实际 gunicorn 的
    post_fork hook 只在 worker 进程里跑。如果用 N>1, 需要用 file lock 或
    gunicorn master/worker 区分; 简单起见, 用环境变量 SM_SYNC_LOOP_ENABLE
    显式开 (生产 unit 不设, 所以默认是关, 留 gunicorn 进程数=1 的安全 case)。
    """
    global _sync_thread
    with _sync_lock:
        if _sync_thread is not None and _sync_thread.is_alive():
            return
        if os.environ.get("SM_SYNC_LOOP_ENABLE", "true").lower() != "true":
            log.info("sync loop disabled (SM_SYNC_LOOP_ENABLE != true)")
            return
        interval = int(os.environ.get("SM_SYNC_INTERVAL", "30"))
        if interval < 5:
            interval = 5  # 5s 下限, 防止在 hot loop 里被滥用
        im = app.instance_manager

        def _loop():
            log.info("SOCKS5 sync loop started, interval=%ds", interval)
            while True:
                try:
                    im.sync_instances()
                except Exception as e:
                    log.exception("sync_instances failed: %s", e)
                time.sleep(interval)

        t = threading.Thread(target=_loop, daemon=True, name="socks5-sync")
        t.start()
        _sync_thread = t
        atexit.register(lambda: log.info("socks5-sync loop exiting (atexit)"))


log = logging.getLogger("socks.app")


def create_app():
    # 确保密码已设置（首次启动自动设置默认密码）
    from auth import get_or_set_password
    get_or_set_password()

    app = Flask(__name__)
    app.config.from_object(Config)
    app.config["SECRET_KEY"] = os.environ.get(
        "SM_SECRET_KEY", "sm-dev-key-change-in-prod"
    )
    app.config["SESSION_COOKIE_HTTPONLY"] = True
    app.config["SESSION_COOKIE_SAMESITE"] = "Lax"

    # 日志
    logging.basicConfig(
        level=getattr(logging, Config.LOG_LEVEL),
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    # 数据库
    db.init_app(app)
    app.db = db

    # 全局服务注册
    with app.app_context():
        db.create_all()
        # 初始化全局 db.app
        from database import db as _db
        _db.app = app

    # 用户服务
    user_service = UserService()
    user_service.set_db_app(app)
    app.user_service = user_service

    # 实例管理器
    instance_manager = InstanceManager(user_service, flask_app=app)
    app.instance_manager = instance_manager

    # 备份目录
    set_backup_dir(Config.BACKUP_DIR)

    # 注册蓝图
    from auth import bp as auth_bp
    from web.routes import bp as web_bp
    from api.v1 import bp as api_bp

    app.register_blueprint(auth_bp)
    app.register_blueprint(web_bp)
    app.register_blueprint(api_bp, url_prefix="/api")

    # 启动后台 sync 线程 (在主线程, 不在 worker 上下文)
    # 必须在 register_blueprint 之后, 调一次 sync_instances 把当前已配的
    # SOCKS5 实例拉起来, 再进 loop 周期探活
    _start_sync_loop(app)
    try:
        instance_manager.sync_instances()
    except Exception as e:
        log.exception("initial sync_instances failed: %s", e)

    return app