Skip to main content

sokuto_server/
lib.rs

1//! # sokuto-server
2//!
3//! Axum をベースとした TypeSafe AI Jev 互換 HTTP API サーバーを提供するクレート。
4//!
5//! ## 概要
6//! - `POST /v1/systemone` 推論エンドポイントの実装。
7//! - OpenAPI 3.1 仕様書のコンパイル時自動導出と Swagger UI (`GET /swagger-ui`) の統合。
8//! - 運用監視 (`GET /health`, `GET /ready`, `GET /metrics`)。
9//! - Tokio 非同期ワーカースレッドのストールを防止するブロッキング分離。
10
11use std::net::SocketAddr;
12use std::sync::Arc;
13
14use axum::Extension;
15use axum::Router;
16use axum::extract::DefaultBodyLimit;
17use axum::routing::{get, post};
18use metrics_exporter_prometheus::PrometheusHandle;
19use tower_http::cors::CorsLayer;
20use tower_http::trace::TraceLayer;
21
22pub use sokuto_core as core;
23pub use sokuto_runtime as runtime;
24
25pub mod error;
26pub mod guardrails;
27pub mod handlers;
28pub mod metrics;
29pub mod openapi;
30pub mod schema;
31pub mod service;
32pub mod state;
33
34pub use error::{ErrorResponse, ServerError};
35pub use guardrails::{GuardrailConfig, GuardrailPipeline};
36pub use openapi::{export_openapi_json, generate_openapi_spec};
37pub use schema::dag::{DagRequest, DagResponse};
38pub use service::{InferenceOutput, InferenceService};
39pub use state::AppState;
40
41/// サーバーのバージョン情報を取得する。
42pub fn version() -> &'static str {
43    env!("CARGO_PKG_VERSION")
44}
45
46/// Jev HTTP API サーバーのルーターを構築する。
47///
48/// ソケットバインドを行わずに `Router` を返却するため、
49/// `tower::ServiceExt::oneshot` 等を用いたインメモリ統合テストが可能である。
50///
51/// # 引数
52/// - `state`: アプリケーション共有状態。
53/// - `prometheus_handle`: Prometheus メトリクスレンダリングハンドル(任意)。
54pub fn create_router(state: Arc<AppState>, prometheus_handle: Option<PrometheusHandle>) -> Router {
55    let api_routes = Router::new()
56        .route("/v1/systemone", post(handlers::system_one_handler))
57        .route(
58            "/v1/systemone/state-batch",
59            post(handlers::state_batch_handler),
60        )
61        .route("/v1/systemone/dag", post(handlers::dag_handler))
62        .route("/health", get(handlers::health_handler))
63        .route("/ready", get(handlers::ready_handler))
64        .route("/metrics", get(handlers::metrics_handler))
65        .layer(Extension(prometheus_handle))
66        .with_state(state);
67
68    let router = Router::new()
69        .merge(api_routes)
70        // 巨大リクエストによる OOM を防止するボディ制限(10MB)
71        .layer(DefaultBodyLimit::max(10 * 1024 * 1024))
72        .layer(CorsLayer::permissive())
73        .layer(TraceLayer::new_for_http());
74
75    // Swagger UI のマウント(feature = "swagger" 有効時のみ実体展開)
76    openapi::mount_swagger_ui(router)
77}
78
79/// 指定されたアドレスで HTTP サーバーを起動する。
80///
81/// グレースフルシャットダウンシグナル(Ctrl+C)を受信するまでイベントループを継続する。
82///
83/// # 引数
84/// - `addr`: バインドするソケットアドレス。
85/// - `state`: アプリケーション共有状態。
86/// - `prometheus_handle`: Prometheus レコーダーハンドル。
87pub async fn run_server(
88    addr: SocketAddr,
89    state: Arc<AppState>,
90    prometheus_handle: Option<PrometheusHandle>,
91) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
92    let app = create_router(state, prometheus_handle);
93
94    tracing::info!("HTTP サーバーを起動します (アドレス: http://{})...", addr);
95    tracing::info!("Swagger UI: http://{}/swagger-ui", addr);
96    tracing::info!("OpenAPI JSON: http://{}/api-docs/openapi.json", addr);
97    tracing::info!("メトリクス: http://{}/metrics", addr);
98    tracing::info!("ヘルスチェック: http://{addr}/health, http://{addr}/ready");
99
100    let listener = tokio::net::TcpListener::bind(addr).await?;
101    axum::serve(listener, app)
102        .with_graceful_shutdown(shutdown_signal())
103        .await?;
104
105    tracing::info!("HTTP サーバーが正常に終了しました。");
106    Ok(())
107}
108
109/// グレースフルシャットダウンシグナル待機用フューチャー。
110async fn shutdown_signal() {
111    let ctrl_c = async {
112        tokio::signal::ctrl_c()
113            .await
114            .expect("Ctrl+C シグナルの捕捉に失敗しました。");
115    };
116
117    #[cfg(unix)]
118    let terminate = async {
119        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
120            .expect("SIGTERM シグナルの捕捉に失敗しました。")
121            .recv()
122            .await;
123    };
124
125    #[cfg(not(unix))]
126    let terminate = std::future::pending::<()>();
127
128    tokio::select! {
129        _ = ctrl_c => {
130            tracing::info!("シャットダウンシグナル (Ctrl+C) を受信しました。");
131        },
132        _ = terminate => {
133            tracing::info!("シャットダウンシグナル (SIGTERM) を受信しました。");
134        },
135    }
136}