"""2 リスナーの起動。 1 コンテナ 1 プロセスのまま、公開ポートと control ポートを同一イベント ループへ並べる。 どちらかのリスナーが終了したらプロセス全体を終了させ、コンテナの再起動に委ねる。 SIGTERM では新規受理を止め、``GRACEFUL_TIMEOUT`` だけ進行中のストリームを待つ。 """ from __future__ import annotations import asyncio import contextlib import logging import signal import sys from typing import Callable, Iterator, Sequence import uvicorn from .config import Settings, load_settings from .control import create_control_app from .db import Database, LogWriter, MaintenanceTask from .keys import KeyCache from .limits import RateLimiter from .logging_utils import configure_logging, register_secret from .proxy import create_proxy_app from .upstream import UpstreamClient logger = logging.getLogger("gateway") class ManagedServer(uvicorn.Server): """シグナル処理を呼び出し側に委ねる uvicorn サーバー。 ``uvicorn.Server.serve`` は既定でシグナル ハンドラを差し替えるため、 2 つ並べると後から起動した側だけが終了要求を受け取ってしまう。 """ @contextlib.contextmanager def capture_signals(self) -> Iterator[None]: # noqa: D102 yield def build_servers( settings: Settings, *, db: Database, keys: KeyCache, limiter: RateLimiter, upstream: UpstreamClient, logs: LogWriter, ) -> tuple[ManagedServer, ManagedServer]: """公開用と control 用のサーバーを組み立てる。""" proxy_app = create_proxy_app( settings=settings, keys=keys, limiter=limiter, upstream=upstream, logs=logs ) # control 側の /healthz から公開リスナーの生死を見るための後方参照。 # 公開リスナーは control アプリより後に組み立てるため、辞書経由で解決する。 listeners: dict[str, ManagedServer] = {} def public_listener_alive() -> bool: server = listeners.get("proxy") return server is not None and bool(server.started) and not server.should_exit control_app = create_control_app( settings=settings, db=db, keys=keys, upstream=upstream, limiter=limiter, listeners_alive=public_listener_alive, ) graceful = int(settings.graceful_timeout) proxy = ManagedServer( uvicorn.Config( proxy_app, host=settings.host, port=settings.port, lifespan="off", access_log=True, server_header=False, date_header=True, timeout_graceful_shutdown=graceful, proxy_headers=False, ) ) control = ManagedServer( uvicorn.Config( control_app, host=settings.host, port=settings.control_port, lifespan="off", access_log=False, server_header=False, timeout_graceful_shutdown=graceful, proxy_headers=False, ) ) listeners["proxy"] = proxy return proxy, control def _install_signal_handlers( servers: Sequence[ManagedServer], graceful_timeout: float, is_finished: Callable[[], bool], ) -> None: """SIGTERM / SIGINT で緩やかに終了する。 新規受理を止めたうえで ``graceful_timeout`` だけ進行中のストリームを待ち、 それでも終わらなければ強制終了へ切り替える。 """ loop = asyncio.get_running_loop() def _shutdown(signame: str) -> None: logger.info("%s を受信しました。新規受理を停止します", signame) for server in servers: server.should_exit = True def _force() -> None: if is_finished(): return logger.warning("待避時間を超えたため進行中の接続を打ち切ります") for server in servers: server.force_exit = True loop.call_later(graceful_timeout, _force) for signame in ("SIGTERM", "SIGINT"): sig = getattr(signal, signame, None) if sig is None: continue with contextlib.suppress(NotImplementedError): loop.add_signal_handler(sig, _shutdown, signame) async def serve_listener(server: ManagedServer, label: str) -> None: """1 つのリスナーを走らせる。 待ち受けを確保できないと uvicorn は ``sys.exit`` を呼ぶ。``SystemExit`` は ``BaseException`` であり、タスクから抜けてイベント ループごと巻き上げるため、 ``run`` の終了処理が走らなくなる。通常の例外へ変換して呼び出し側で扱う。 """ try: await server.serve() except SystemExit as exc: # pragma: no cover - 待ち受け確保の失敗時のみ raise RuntimeError(f"{label} リスナーを起動できませんでした（終了コード {exc.code}）") from exc async def run(settings: Settings | None = None) -> int: """ゲートウェイを起動し、終了コードを返す。""" settings = settings or load_settings() configure_logging() register_secret(settings.upstream_api_key, settings.control_token) db = Database(settings.db_path) await db.connect() keys = KeyCache(db=db, pepper=settings.key_pepper, ttl=settings.key_cache_ttl) await keys.start() limiter = RateLimiter(settings) await limiter.restore(db) upstream = UpstreamClient(settings) logs = LogWriter(db) await logs.start() maintenance = MaintenanceTask(db=db, settings=settings) await maintenance.start() snapshot_task = asyncio.create_task( limiter.snapshot_loop(db, settings.rate_snapshot_interval), name="rate-snapshot" ) proxy, control = build_servers( settings, db=db, keys=keys, limiter=limiter, upstream=upstream, logs=logs ) servers = (proxy, control) finished = False _install_signal_handlers(servers, settings.graceful_timeout, lambda: finished) logger.info( "公開ポート %s:%s / control ポート %s:%s で待ち受けます（パス方針: %s）", settings.host, settings.port, settings.host, settings.control_port, settings.path_policy, ) tasks = { asyncio.create_task(serve_listener(proxy, "公開"), name="proxy-server"): "公開", asyncio.create_task(serve_listener(control, "control"), name="control-server"): "control", } exit_code = 0 try: done, pending = await asyncio.wait(tasks.keys(), return_when=asyncio.FIRST_COMPLETED) for task in done: label = tasks[task] error = task.exception() if error is not None: logger.error("%s リスナーが異常終了しました: %s", label, error) exit_code = 1 else: logger.info("%s リスナーが終了しました", label) # 片方が落ちた時点でプロセス全体を終了させ、再起動をコンテナに委ねる。 for server in servers: server.should_exit = True if pending: _, still_running = await asyncio.wait(pending, timeout=settings.graceful_timeout) if still_running: for server in servers: server.force_exit = True await asyncio.wait(still_running, timeout=5.0) for task in still_running: task.cancel() finally: finished = True snapshot_task.cancel() with contextlib.suppress(asyncio.CancelledError): await snapshot_task with contextlib.suppress(Exception): await limiter.snapshot(db) await maintenance.stop() await keys.stop() await logs.stop() await upstream.aclose() await db.close() logger.info("終了処理を完了しました") return exit_code def main() -> int: try: return asyncio.run(run()) except KeyboardInterrupt: # pragma: no cover return 0 except SystemExit as exc: # 待ち受けの確保に失敗すると uvicorn が sys.exit を呼ぶ。 # 例外の投げ上げではなく終了コードとして扱い、再起動をコンテナに委ねる。 code = exc.code if isinstance(exc.code, int) else 1 logger.error("リスナーを起動できなかったため終了します（終了コード %s）", code) return code if __name__ == "__main__": # pragma: no cover sys.exit(main())