893 lines
31 KiB
Rust
893 lines
31 KiB
Rust
#![recursion_limit = "256"]
|
|
|
|
mod admin;
|
|
mod admin_accounts;
|
|
mod admin_recharge;
|
|
mod ai_tasks;
|
|
mod aliyun_matting;
|
|
mod api_response;
|
|
mod app;
|
|
mod asset_billing;
|
|
mod assets;
|
|
mod auth;
|
|
mod auth_me;
|
|
mod auth_payload;
|
|
mod auth_public_user;
|
|
mod auth_session;
|
|
mod auth_sessions;
|
|
mod backpressure;
|
|
mod bgfilter_worker;
|
|
mod character_animation_assets;
|
|
mod character_visual_assets;
|
|
mod config;
|
|
mod custom_world_asset_prompts;
|
|
mod editor_agent;
|
|
mod editor_background_music_prompt_assist;
|
|
mod editor_generation_config;
|
|
mod editor_generation_queue;
|
|
mod editor_green_screen;
|
|
mod editor_project;
|
|
mod editor_project_icon;
|
|
mod editor_screen_background_decision;
|
|
mod editor_screen_background_filter;
|
|
mod editor_sound_effect_prompt_assist;
|
|
mod error_middleware;
|
|
mod external_api_audit;
|
|
mod external_api_auth;
|
|
mod external_api_keys;
|
|
mod external_assets_api;
|
|
mod external_editor_api;
|
|
mod external_generation;
|
|
mod external_generation_worker;
|
|
mod external_generation_worker_controller;
|
|
mod external_mcp;
|
|
mod external_skill_api;
|
|
mod frontend_runtime_config;
|
|
mod generated_image_assets;
|
|
mod health;
|
|
mod http_error;
|
|
mod hyper3d_generation;
|
|
mod llm;
|
|
mod llm_model_routing;
|
|
#[cfg(test)]
|
|
mod llm_prompt_test_support;
|
|
mod login_options;
|
|
mod logout;
|
|
mod logout_all;
|
|
mod modules;
|
|
mod openai_image_generation;
|
|
mod password_entry;
|
|
mod password_management;
|
|
mod phone_auth;
|
|
mod platform_errors;
|
|
mod process_metrics;
|
|
mod profile_identity;
|
|
mod profile_recharge_expiration_listener;
|
|
mod profile_recharge_refund_reconciliation;
|
|
mod prompt;
|
|
mod refresh_session;
|
|
mod registration_reward;
|
|
mod request_context;
|
|
mod response_headers;
|
|
mod runtime_profile;
|
|
mod runtime_settings;
|
|
mod session_client;
|
|
mod state;
|
|
mod telemetry;
|
|
mod tracking;
|
|
mod tracking_outbox;
|
|
mod vector_engine_audio_generation;
|
|
mod volcengine_speech;
|
|
mod wallet_refund_outbox;
|
|
mod wechat;
|
|
mod work_author;
|
|
|
|
use shared_logging::{OtelConfig, init_tracing};
|
|
use socket2::{Domain, Protocol, Socket, Type};
|
|
use std::{
|
|
collections::HashSet,
|
|
env, fs, future, io,
|
|
net::{SocketAddr, TcpListener as StdTcpListener},
|
|
panic,
|
|
sync::Arc,
|
|
thread,
|
|
time::Duration,
|
|
};
|
|
use tokio::net::TcpListener;
|
|
use tokio::runtime::Builder as TokioRuntimeBuilder;
|
|
use tokio::time::timeout;
|
|
use tracing::{error, info, warn};
|
|
|
|
use crate::{
|
|
app::{build_router, build_spacetime_unavailable_router},
|
|
bgfilter_worker::{build_bgfilter_worker_router, validate_bgfilter_internal_token},
|
|
config::{AppConfig, ProcessRole},
|
|
external_generation_worker::run_external_generation_worker,
|
|
external_generation_worker_controller::run_external_generation_worker_controller,
|
|
profile_recharge_expiration_listener::spawn_profile_recharge_expiration_listener,
|
|
profile_recharge_refund_reconciliation::spawn_profile_recharge_refund_reconciliation_worker,
|
|
state::{AppState, AppStateInitError},
|
|
tracking_outbox::TrackingOutbox,
|
|
wallet_refund_outbox::WalletRefundOutbox,
|
|
};
|
|
|
|
const API_SERVER_STARTUP_STACK_SIZE_BYTES: usize = 32 * 1024 * 1024;
|
|
const AUTH_STORE_STARTUP_RESTORE_TIMEOUT: Duration = Duration::from_secs(8);
|
|
const AUTH_STORE_STARTUP_RETRY_INTERVAL: Duration = Duration::from_secs(5);
|
|
|
|
#[derive(Clone)]
|
|
struct ShutdownContext {
|
|
app_state: Option<AppState>,
|
|
tracking_outbox: Option<Arc<TrackingOutbox>>,
|
|
wallet_refund_outbox: Option<Arc<WalletRefundOutbox>>,
|
|
outbox_flush_timeout: Duration,
|
|
}
|
|
|
|
fn main() -> Result<(), io::Error> {
|
|
// Windows 本地调试下 Axum 路由树和启动恢复链较重，显式放大启动线程栈，避免 debug 构建在进入监听前栈溢出。
|
|
let server_thread = thread::Builder::new()
|
|
.name("api-server-bootstrap".to_string())
|
|
.stack_size(API_SERVER_STARTUP_STACK_SIZE_BYTES)
|
|
.spawn(|| {
|
|
load_local_env_files();
|
|
let config = AppConfig::from_env();
|
|
let mut runtime_builder = TokioRuntimeBuilder::new_multi_thread();
|
|
runtime_builder
|
|
.enable_all()
|
|
.thread_name("api-server-worker")
|
|
.thread_stack_size(API_SERVER_STARTUP_STACK_SIZE_BYTES);
|
|
if let Some(worker_threads) = config.worker_threads {
|
|
runtime_builder.worker_threads(worker_threads);
|
|
}
|
|
|
|
runtime_builder.build()?.block_on(run_server(config))
|
|
})?;
|
|
|
|
match server_thread.join() {
|
|
Ok(result) => result,
|
|
Err(payload) => panic::resume_unwind(payload),
|
|
}
|
|
}
|
|
|
|
async fn run_server(config: AppConfig) -> Result<(), io::Error> {
|
|
validate_bgfilter_internal_token_for_startup(&config).map_err(io::Error::other)?;
|
|
init_tracing(
|
|
&config.log_filter,
|
|
OtelConfig {
|
|
enabled: config.otel_enabled,
|
|
},
|
|
)?;
|
|
process_metrics::register_process_metrics();
|
|
telemetry::register_http_runtime_metrics();
|
|
|
|
if config.process_role.runs_bgfilter_worker() {
|
|
return run_bgfilter_worker_role(config).await;
|
|
}
|
|
|
|
if !config.process_role.runs_http() {
|
|
return run_worker_only(config).await;
|
|
}
|
|
|
|
run_http_role(config).await
|
|
}
|
|
|
|
fn validate_bgfilter_internal_token_for_startup(config: &AppConfig) -> Result<(), String> {
|
|
if should_validate_bgfilter_internal_token_for_startup(config.process_role) {
|
|
validate_bgfilter_internal_token(config.bgfilter_internal_token.as_deref())?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn should_validate_bgfilter_internal_token_for_startup(process_role: ProcessRole) -> bool {
|
|
matches!(
|
|
process_role,
|
|
ProcessRole::Api
|
|
| ProcessRole::BgfilterWorker
|
|
| ProcessRole::ExternalGenerationWorker
|
|
| ProcessRole::All
|
|
)
|
|
}
|
|
|
|
async fn run_bgfilter_worker_role(mut config: AppConfig) -> Result<(), io::Error> {
|
|
let (concurrency, single_image_estimate_ms, max_requests) =
|
|
required_bgfilter_worker_capacity_from_env()?;
|
|
config.bgfilter_worker_concurrency = concurrency;
|
|
config.editor_bgfilter_single_image_estimate_ms = single_image_estimate_ms;
|
|
config.bgfilter_worker_max_requests = max_requests;
|
|
let bind_address = format!(
|
|
"{}:{}",
|
|
config.bgfilter_worker_host, config.bgfilter_worker_port
|
|
)
|
|
.parse::<SocketAddr>()
|
|
.map_err(|error| io::Error::other(format!("bgfilter-worker 监听地址无效：{error}")))?;
|
|
if !bind_address.ip().is_loopback() {
|
|
return Err(io::Error::other(format!(
|
|
"bgfilter-worker 首版只允许监听 loopback，当前地址为 {bind_address}"
|
|
)));
|
|
}
|
|
let listen_backlog = config.listen_backlog;
|
|
let outbox_flush_timeout = config.shutdown_outbox_flush_timeout;
|
|
let listener = build_tcp_listener(bind_address, listen_backlog)?;
|
|
|
|
configure_bgfilter_worker_outboxes(&mut config);
|
|
let state = AppState::new_with_empty_auth_store(config)
|
|
.map_err(|error| io::Error::other(format!("初始化 bgfilter-worker 状态失败：{error}")))?;
|
|
let (router, task_tracker) = build_bgfilter_worker_router(state.clone())
|
|
.map_err(|error| io::Error::other(format!("初始化 bgfilter-worker 路由失败：{error}")))?;
|
|
let tracking_outbox = state.tracking_outbox();
|
|
if let Some(outbox) = tracking_outbox.clone() {
|
|
outbox.spawn_worker();
|
|
}
|
|
let shutdown_context = ShutdownContext {
|
|
app_state: Some(state),
|
|
tracking_outbox,
|
|
wallet_refund_outbox: None,
|
|
outbox_flush_timeout,
|
|
};
|
|
info!(
|
|
%bind_address,
|
|
listen_backlog,
|
|
process_role = ProcessRole::BgfilterWorker.as_str(),
|
|
"bgfilter-worker 已开始监听内部 HTTP"
|
|
);
|
|
let shutdown_tracker = task_tracker.clone();
|
|
let shutdown_context_for_signal = shutdown_context.clone();
|
|
let result = axum::serve(listener, router)
|
|
.with_graceful_shutdown(async move {
|
|
shutdown_signal(shutdown_context_for_signal).await;
|
|
shutdown_tracker.close();
|
|
})
|
|
.await;
|
|
task_tracker.close();
|
|
task_tracker.wait_for_drain().await;
|
|
finalize_shutdown(shutdown_context).await;
|
|
result
|
|
}
|
|
|
|
fn configure_bgfilter_worker_outboxes(config: &mut AppConfig) {
|
|
// 多进程不能操作同一个 active 文件；worker 从共享基础目录派生自己的持久子目录。
|
|
config.tracking_outbox_enabled = true;
|
|
config.tracking_outbox_dir = config.tracking_outbox_dir.join("bgfilter-worker");
|
|
config.wallet_refund_outbox_enabled = false;
|
|
}
|
|
|
|
const DEFAULT_BGFILTER_WORKER_MAX_REQUESTS_FUSE: usize = 2_048;
|
|
|
|
fn required_bgfilter_worker_capacity_from_env() -> Result<(usize, u64, usize), io::Error> {
|
|
let concurrency = env::var("GENARRATIVE_BGFILTER_WORKER_CONCURRENCY").ok();
|
|
let single_image_estimate_ms =
|
|
env::var("GENARRATIVE_EDITOR_BGFILTER_SINGLE_IMAGE_ESTIMATE_MS").ok();
|
|
let max_requests = env::var("GENARRATIVE_BGFILTER_WORKER_MAX_REQUESTS").ok();
|
|
parse_required_bgfilter_worker_capacity(
|
|
concurrency.as_deref(),
|
|
single_image_estimate_ms.as_deref(),
|
|
max_requests.as_deref(),
|
|
)
|
|
}
|
|
|
|
fn parse_required_bgfilter_worker_capacity(
|
|
concurrency: Option<&str>,
|
|
single_image_estimate_ms: Option<&str>,
|
|
max_requests: Option<&str>,
|
|
) -> Result<(usize, u64, usize), io::Error> {
|
|
fn parse_required_positive(name: &str, raw: Option<&str>) -> Result<usize, io::Error> {
|
|
let value = raw
|
|
.map(strip_env_value)
|
|
.map(|value| value.trim().to_string())
|
|
.filter(|value| !value.is_empty())
|
|
.ok_or_else(|| io::Error::other(format!("bgfilter-worker 启动必须显式配置 {name}")))?;
|
|
let parsed = value.parse::<usize>().map_err(|error| {
|
|
io::Error::other(format!(
|
|
"bgfilter-worker 配置 {name} 不是有效正整数：{error}"
|
|
))
|
|
})?;
|
|
if parsed == 0 {
|
|
return Err(io::Error::other(format!(
|
|
"bgfilter-worker 配置 {name} 必须大于 0"
|
|
)));
|
|
}
|
|
Ok(parsed)
|
|
}
|
|
|
|
let concurrency =
|
|
parse_required_positive("GENARRATIVE_BGFILTER_WORKER_CONCURRENCY", concurrency)?;
|
|
let single_image_estimate_ms = parse_required_positive(
|
|
"GENARRATIVE_EDITOR_BGFILTER_SINGLE_IMAGE_ESTIMATE_MS",
|
|
single_image_estimate_ms,
|
|
)? as u64;
|
|
// Q 已降级为 admission 保险丝：可缺省（默认 2048），显式配置时仍必须为正且不小于 N。
|
|
let max_requests = match max_requests
|
|
.map(strip_env_value)
|
|
.map(|value| value.trim().to_string())
|
|
.filter(|value| !value.is_empty())
|
|
{
|
|
None => DEFAULT_BGFILTER_WORKER_MAX_REQUESTS_FUSE,
|
|
Some(raw) => {
|
|
let parsed = raw.parse::<usize>().map_err(|error| {
|
|
io::Error::other(format!(
|
|
"bgfilter-worker 配置 GENARRATIVE_BGFILTER_WORKER_MAX_REQUESTS 不是有效正整数：{error}"
|
|
))
|
|
})?;
|
|
if parsed == 0 {
|
|
return Err(io::Error::other(
|
|
"bgfilter-worker 配置 GENARRATIVE_BGFILTER_WORKER_MAX_REQUESTS 必须大于 0",
|
|
));
|
|
}
|
|
parsed
|
|
}
|
|
};
|
|
if max_requests < concurrency {
|
|
return Err(io::Error::other(
|
|
"GENARRATIVE_BGFILTER_WORKER_MAX_REQUESTS 不能小于 GENARRATIVE_BGFILTER_WORKER_CONCURRENCY",
|
|
));
|
|
}
|
|
Ok((concurrency, single_image_estimate_ms, max_requests))
|
|
}
|
|
|
|
async fn run_worker_only(config: AppConfig) -> Result<(), io::Error> {
|
|
let process_role = config.process_role;
|
|
let state = build_non_http_app_state_for_startup(config).map_err(|error| {
|
|
io::Error::other(format!(
|
|
"初始化 external generation worker 状态失败：{error}"
|
|
))
|
|
})?;
|
|
state
|
|
.spacetime_client()
|
|
.get_external_generation_queue_stats()
|
|
.await
|
|
.map_err(|error| {
|
|
io::Error::other(format!("验证外部生成队列运行时服务身份失败：{error}"))
|
|
})?;
|
|
spawn_common_app_state_background_workers(&state);
|
|
info!(
|
|
process_role = process_role.as_str(),
|
|
"api-server 以非 HTTP 角色启动"
|
|
);
|
|
if process_role.runs_external_generation_worker() {
|
|
run_external_generation_worker(state).await
|
|
} else if process_role.runs_external_generation_controller() {
|
|
run_external_generation_worker_controller(state).await
|
|
} else {
|
|
Err(io::Error::other(format!(
|
|
"不支持的非 HTTP 进程角色：{}",
|
|
process_role.as_str()
|
|
)))
|
|
}
|
|
}
|
|
|
|
fn build_non_http_app_state_for_startup(
|
|
config: AppConfig,
|
|
) -> Result<AppState, state::AppStateInitError> {
|
|
let process_role = config.process_role;
|
|
debug_assert!(!should_restore_auth_store_for_startup(process_role));
|
|
info!(
|
|
process_role = process_role.as_str(),
|
|
"非 HTTP 进程跳过 SpacetimeDB 认证投影恢复"
|
|
);
|
|
AppState::new_with_empty_auth_store(config)
|
|
}
|
|
|
|
fn should_restore_auth_store_for_startup(process_role: ProcessRole) -> bool {
|
|
process_role.runs_http()
|
|
}
|
|
|
|
fn should_initialize_editor_generation_pricing_for_startup(process_role: ProcessRole) -> bool {
|
|
process_role.runs_http()
|
|
}
|
|
|
|
async fn run_http_role(config: AppConfig) -> Result<(), io::Error> {
|
|
let bind_address = config.bind_socket_addr();
|
|
let listen_backlog = config.listen_backlog;
|
|
let worker_threads = config.worker_threads;
|
|
let otel_enabled = config.otel_enabled;
|
|
let process_role = config.process_role;
|
|
let outbox_flush_timeout = config.shutdown_outbox_flush_timeout;
|
|
let listener = build_tcp_listener(bind_address, listen_backlog)?;
|
|
|
|
let (router, shutdown_context, worker_state) = match restore_app_state_for_startup(config).await
|
|
{
|
|
Ok(state) => {
|
|
spawn_http_app_state_background_workers(&state, process_role);
|
|
let tracking_outbox = state.tracking_outbox();
|
|
let wallet_refund_outbox = state.wallet_refund_outbox();
|
|
let worker_state = process_role
|
|
.runs_external_generation_worker()
|
|
.then(|| state.clone());
|
|
(
|
|
build_router(state.clone()),
|
|
ShutdownContext {
|
|
app_state: Some(state),
|
|
tracking_outbox,
|
|
wallet_refund_outbox,
|
|
outbox_flush_timeout,
|
|
},
|
|
worker_state,
|
|
)
|
|
}
|
|
Err(AppStateInitError::DependencyUnavailable(message)) => (
|
|
build_spacetime_unavailable_router(message),
|
|
ShutdownContext {
|
|
app_state: None,
|
|
tracking_outbox: None,
|
|
wallet_refund_outbox: None,
|
|
outbox_flush_timeout,
|
|
},
|
|
None,
|
|
),
|
|
Err(error) => {
|
|
return Err(std::io::Error::other(format!(
|
|
"初始化应用状态失败：{error}"
|
|
)));
|
|
}
|
|
};
|
|
|
|
info!(
|
|
%bind_address,
|
|
listen_backlog,
|
|
worker_threads = worker_threads.unwrap_or(0),
|
|
otel_enabled,
|
|
process_role = process_role.as_str(),
|
|
"api-server 已完成 tracing 初始化并开始监听"
|
|
);
|
|
|
|
let http_server = axum::serve(listener, router)
|
|
.with_graceful_shutdown(shutdown_signal(shutdown_context.clone()));
|
|
let result = if let Some(worker_state) = worker_state {
|
|
tokio::select! {
|
|
result = http_server => result,
|
|
result = run_external_generation_worker(worker_state) => result,
|
|
}
|
|
} else {
|
|
http_server.await
|
|
};
|
|
finalize_shutdown(shutdown_context).await;
|
|
result
|
|
}
|
|
|
|
async fn shutdown_signal(context: ShutdownContext) {
|
|
let signal = wait_for_shutdown_signal().await;
|
|
if let Some(state) = context.app_state.as_ref() {
|
|
state.mark_not_ready();
|
|
}
|
|
info!(
|
|
signal,
|
|
"api-server 收到退出信号，已标记 readiness 不可用并开始排空 HTTP 请求"
|
|
);
|
|
}
|
|
|
|
async fn wait_for_shutdown_signal() -> &'static str {
|
|
#[cfg(unix)]
|
|
{
|
|
tokio::select! {
|
|
signal = wait_for_ctrl_c_signal() => signal,
|
|
signal = wait_for_sigterm_signal() => signal,
|
|
}
|
|
}
|
|
|
|
#[cfg(not(unix))]
|
|
{
|
|
wait_for_ctrl_c_signal().await
|
|
}
|
|
}
|
|
|
|
async fn wait_for_ctrl_c_signal() -> &'static str {
|
|
if let Err(error) = tokio::signal::ctrl_c().await {
|
|
error!(error = %error, "监听 SIGINT 失败，无法通过 Ctrl-C 触发优雅退出");
|
|
future::pending::<()>().await;
|
|
}
|
|
"sigint"
|
|
}
|
|
|
|
#[cfg(unix)]
|
|
async fn wait_for_sigterm_signal() -> &'static str {
|
|
let mut signal = match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
|
|
{
|
|
Ok(signal) => signal,
|
|
Err(error) => {
|
|
error!(error = %error, "监听 SIGTERM 失败，无法通过 systemd terminate 触发优雅退出");
|
|
future::pending::<()>().await;
|
|
unreachable!("pending future never returns");
|
|
}
|
|
};
|
|
signal.recv().await;
|
|
"sigterm"
|
|
}
|
|
|
|
async fn finalize_shutdown(context: ShutdownContext) {
|
|
if let Some(state) = context.app_state.as_ref() {
|
|
state.mark_not_ready();
|
|
}
|
|
|
|
if context.outbox_flush_timeout.is_zero() {
|
|
warn!("api-server 退出时 outbox flush timeout 为 0，跳过主动 flush");
|
|
return;
|
|
}
|
|
|
|
let timeout_ms = context
|
|
.outbox_flush_timeout
|
|
.as_millis()
|
|
.min(u128::from(u64::MAX)) as u64;
|
|
if let Some(outbox) = context.tracking_outbox {
|
|
info!(timeout_ms, "api-server 退出前封存并 flush tracking outbox");
|
|
match timeout(context.outbox_flush_timeout, outbox.flush_for_shutdown()).await {
|
|
Ok(Ok(())) => {
|
|
info!("api-server 退出前 tracking outbox flush 完成");
|
|
}
|
|
Ok(Err(error)) => {
|
|
warn!(
|
|
error = %error,
|
|
"api-server 退出前 tracking outbox flush 未完成，已保留本地文件等待下次启动重试"
|
|
);
|
|
}
|
|
Err(_) => {
|
|
warn!(
|
|
timeout_ms,
|
|
"api-server 退出前 tracking outbox flush 超时，已保留本地文件等待下次启动重试"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
if let Some(outbox) = context.wallet_refund_outbox {
|
|
info!(timeout_ms, "api-server 退出前 flush wallet refund outbox");
|
|
match timeout(context.outbox_flush_timeout, outbox.flush_for_shutdown()).await {
|
|
Ok(Ok(())) => {
|
|
info!("api-server 退出前 wallet refund outbox flush 完成");
|
|
}
|
|
Ok(Err(error)) => {
|
|
warn!(
|
|
error = %error,
|
|
"api-server 退出前 wallet refund outbox flush 未完成，已保留本地文件等待下次启动重试"
|
|
);
|
|
}
|
|
Err(_) => {
|
|
warn!(
|
|
timeout_ms,
|
|
"api-server 退出前 wallet refund outbox flush 超时，已保留本地文件等待下次启动重试"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn spawn_common_app_state_background_workers(state: &AppState) {
|
|
if let Some(outbox) = state.tracking_outbox() {
|
|
outbox.spawn_worker();
|
|
}
|
|
if let Some(outbox) = state.wallet_refund_outbox() {
|
|
outbox.spawn_worker();
|
|
}
|
|
}
|
|
|
|
fn spawn_http_app_state_background_workers(state: &AppState, process_role: ProcessRole) {
|
|
spawn_common_app_state_background_workers(state);
|
|
if should_start_profile_recharge_expiration_listener(process_role) {
|
|
spawn_profile_recharge_expiration_listener(state.clone());
|
|
spawn_profile_recharge_refund_reconciliation_worker(state.clone());
|
|
}
|
|
}
|
|
|
|
fn should_start_profile_recharge_expiration_listener(process_role: ProcessRole) -> bool {
|
|
process_role.runs_http()
|
|
}
|
|
|
|
fn build_tcp_listener(
|
|
bind_address: SocketAddr,
|
|
listen_backlog: i32,
|
|
) -> Result<TcpListener, io::Error> {
|
|
let domain = Domain::for_address(bind_address);
|
|
let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))?;
|
|
socket.set_reuse_address(true)?;
|
|
socket.set_nonblocking(true)?;
|
|
socket.bind(&bind_address.into())?;
|
|
socket.listen(listen_backlog)?;
|
|
TcpListener::from_std(StdTcpListener::from(socket))
|
|
}
|
|
|
|
async fn restore_app_state_for_startup(
|
|
config: AppConfig,
|
|
) -> Result<AppState, state::AppStateInitError> {
|
|
loop {
|
|
match try_restore_app_state_for_startup(config.clone()).await {
|
|
Ok(state) => return Ok(state),
|
|
Err(state::AppStateInitError::DependencyUnavailable(message)) => {
|
|
warn!(
|
|
retry_after_seconds = AUTH_STORE_STARTUP_RETRY_INTERVAL.as_secs(),
|
|
error = %message,
|
|
"启动恢复 SpacetimeDB 认证投影暂不可用，api-server 将继续重试"
|
|
);
|
|
tokio::time::sleep(AUTH_STORE_STARTUP_RETRY_INTERVAL).await;
|
|
}
|
|
Err(error) => return Err(error),
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn try_restore_app_state_for_startup(
|
|
config: AppConfig,
|
|
) -> Result<AppState, state::AppStateInitError> {
|
|
let process_role = config.process_role;
|
|
let state = match timeout(
|
|
AUTH_STORE_STARTUP_RESTORE_TIMEOUT,
|
|
AppState::try_restore_auth_store_from_spacetime(config),
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => result?,
|
|
Err(_) => {
|
|
error!(
|
|
timeout_seconds = AUTH_STORE_STARTUP_RESTORE_TIMEOUT.as_secs(),
|
|
"启动等待 SpacetimeDB 恢复认证投影超时"
|
|
);
|
|
return Err(state::AppStateInitError::DependencyUnavailable(
|
|
"SpacetimeDB 启动恢复认证投影超时".to_string(),
|
|
));
|
|
}
|
|
};
|
|
if should_initialize_editor_generation_pricing_for_startup(process_role) {
|
|
state
|
|
.ensure_editor_generation_runtime_service_identity()
|
|
.await
|
|
.map_err(|error| {
|
|
state::AppStateInitError::DependencyUnavailable(format!(
|
|
"初始化模型定价服务身份失败：{error}"
|
|
))
|
|
})?;
|
|
}
|
|
Ok(state)
|
|
}
|
|
|
|
fn load_local_env_files() {
|
|
let shell_env_keys = protected_env_keys_from(env::vars());
|
|
|
|
for path in [".env", ".env.local", ".env.secrets.local"] {
|
|
load_env_file(path, &shell_env_keys);
|
|
}
|
|
}
|
|
|
|
fn protected_env_keys_from(vars: impl IntoIterator<Item = (String, String)>) -> HashSet<String> {
|
|
vars.into_iter()
|
|
.filter_map(|(key, value)| {
|
|
if value.trim().is_empty() {
|
|
None
|
|
} else {
|
|
Some(key)
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn load_env_file(path: &str, shell_env_keys: &HashSet<String>) {
|
|
let Ok(raw_text) = fs::read_to_string(path) else {
|
|
return;
|
|
};
|
|
|
|
let raw_text = raw_text.trim_start_matches('\u{feff}');
|
|
|
|
for raw_line in raw_text.split('\n') {
|
|
let line = raw_line.trim();
|
|
if line.is_empty() || line.starts_with('#') {
|
|
continue;
|
|
}
|
|
|
|
let Some((raw_key, raw_value)) = line.split_once('=') else {
|
|
continue;
|
|
};
|
|
let key = raw_key.trim().trim_start_matches('\u{feff}');
|
|
if !is_valid_env_key(key) || shell_env_keys.contains(key) {
|
|
continue;
|
|
}
|
|
|
|
// 这里只在启动前、Tokio runtime 创建前写入进程环境，避免并发读写 env。
|
|
unsafe {
|
|
env::set_var(key, strip_env_value(raw_value));
|
|
}
|
|
}
|
|
}
|
|
|
|
fn strip_env_value(raw_value: &str) -> String {
|
|
let value = raw_value.trim_end_matches('\r');
|
|
if value.len() >= 2 {
|
|
let bytes = value.as_bytes();
|
|
let first = bytes[0];
|
|
let last = bytes[value.len() - 1];
|
|
if (first == b'"' && last == b'"') || (first == b'\'' && last == b'\'') {
|
|
return value[1..value.len() - 1].to_string();
|
|
}
|
|
}
|
|
|
|
value.to_string()
|
|
}
|
|
|
|
fn is_valid_env_key(key: &str) -> bool {
|
|
let mut chars = key.chars();
|
|
match chars.next() {
|
|
Some(first) if first == '_' || first.is_ascii_alphabetic() => {}
|
|
_ => return false,
|
|
}
|
|
|
|
chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{
|
|
AUTH_STORE_STARTUP_RETRY_INTERVAL, configure_bgfilter_worker_outboxes, is_valid_env_key,
|
|
parse_required_bgfilter_worker_capacity, protected_env_keys_from,
|
|
should_initialize_editor_generation_pricing_for_startup,
|
|
should_restore_auth_store_for_startup, should_start_profile_recharge_expiration_listener,
|
|
should_validate_bgfilter_internal_token_for_startup, strip_env_value,
|
|
validate_bgfilter_internal_token_for_startup,
|
|
};
|
|
use crate::config::{AppConfig, ProcessRole};
|
|
|
|
#[test]
|
|
fn strip_env_value_removes_wrapping_quotes() {
|
|
assert_eq!(strip_env_value("\"true\""), "true");
|
|
assert_eq!(strip_env_value("'aliyun'"), "aliyun");
|
|
assert_eq!(strip_env_value("plain\r"), "plain");
|
|
}
|
|
|
|
#[test]
|
|
fn bgfilter_worker_capacity_must_be_explicit_positive_and_bounded_by_q() {
|
|
assert_eq!(
|
|
parse_required_bgfilter_worker_capacity(Some("'16'"), Some("5000"), Some(" 128 "))
|
|
.expect("valid explicit N/est/Q"),
|
|
(16, 5_000, 128)
|
|
);
|
|
// Q 是可选保险丝：缺省时取默认值 2048，N 与 est 仍必须显式且为正。
|
|
assert_eq!(
|
|
parse_required_bgfilter_worker_capacity(Some("16"), Some("5000"), None)
|
|
.expect("missing Q falls back to fuse default"),
|
|
(16, 5_000, 2_048)
|
|
);
|
|
for (concurrency, estimate, max_requests) in [
|
|
(None, Some("5000"), Some("128")),
|
|
(Some("16"), None, Some("128")),
|
|
(Some(""), Some("5000"), Some("128")),
|
|
(Some("0"), Some("5000"), Some("128")),
|
|
(Some("16"), Some("0"), Some("128")),
|
|
(Some("four"), Some("5000"), Some("128")),
|
|
(Some("8"), Some("5000"), Some("4")),
|
|
(Some("8"), Some("5000"), Some("0")),
|
|
] {
|
|
assert!(
|
|
parse_required_bgfilter_worker_capacity(concurrency, estimate, max_requests)
|
|
.is_err(),
|
|
"invalid N/est/Q should fail closed: N={concurrency:?}, est={estimate:?}, Q={max_requests:?}"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn bgfilter_worker_uses_its_own_tracking_outbox_directory() {
|
|
let mut config = AppConfig::default();
|
|
let base_dir = config.tracking_outbox_dir.clone();
|
|
config.tracking_outbox_enabled = false;
|
|
config.wallet_refund_outbox_enabled = true;
|
|
|
|
configure_bgfilter_worker_outboxes(&mut config);
|
|
|
|
assert!(config.tracking_outbox_enabled);
|
|
assert_eq!(config.tracking_outbox_dir, base_dir.join("bgfilter-worker"));
|
|
assert!(!config.wallet_refund_outbox_enabled);
|
|
}
|
|
|
|
#[test]
|
|
fn load_env_key_can_strip_utf8_bom_prefix() {
|
|
let key = "\u{feff}SMS_AUTH_ENABLED"
|
|
.trim()
|
|
.trim_start_matches('\u{feff}');
|
|
|
|
assert_eq!(key, "SMS_AUTH_ENABLED");
|
|
}
|
|
|
|
#[test]
|
|
fn is_valid_env_key_accepts_dotenv_key_subset() {
|
|
assert!(is_valid_env_key("SMS_AUTH_ENABLED"));
|
|
assert!(is_valid_env_key("_LOCAL_KEY_1"));
|
|
assert!(!is_valid_env_key("1_BAD"));
|
|
assert!(!is_valid_env_key("BAD-KEY"));
|
|
}
|
|
|
|
#[test]
|
|
fn empty_shell_env_does_not_protect_dotenv_value() {
|
|
let protected = protected_env_keys_from([
|
|
("ALIYUN_OSS_BUCKET".to_string(), "".to_string()),
|
|
("ALIYUN_OSS_ENDPOINT".to_string(), " ".to_string()),
|
|
(
|
|
"ALIYUN_OSS_ACCESS_KEY_ID".to_string(),
|
|
"configured".to_string(),
|
|
),
|
|
]);
|
|
|
|
assert!(!protected.contains("ALIYUN_OSS_BUCKET"));
|
|
assert!(!protected.contains("ALIYUN_OSS_ENDPOINT"));
|
|
assert!(protected.contains("ALIYUN_OSS_ACCESS_KEY_ID"));
|
|
}
|
|
|
|
#[test]
|
|
fn startup_dependency_retry_interval_is_short_enough_for_service_recovery() {
|
|
assert_eq!(AUTH_STORE_STARTUP_RETRY_INTERVAL.as_secs(), 5);
|
|
}
|
|
|
|
#[test]
|
|
fn bgfilter_internal_token_startup_validation_is_limited_to_consumers() {
|
|
for role in [
|
|
ProcessRole::Api,
|
|
ProcessRole::BgfilterWorker,
|
|
ProcessRole::ExternalGenerationWorker,
|
|
ProcessRole::All,
|
|
] {
|
|
assert!(should_validate_bgfilter_internal_token_for_startup(role));
|
|
let mut config = AppConfig::default();
|
|
config.process_role = role;
|
|
config.bgfilter_internal_token = Some("invalid token".to_string());
|
|
assert!(validate_bgfilter_internal_token_for_startup(&config).is_err());
|
|
}
|
|
assert!(!should_validate_bgfilter_internal_token_for_startup(
|
|
ProcessRole::ExternalGenerationController
|
|
));
|
|
let mut controller_config = AppConfig::default();
|
|
controller_config.process_role = ProcessRole::ExternalGenerationController;
|
|
controller_config.bgfilter_internal_token = Some("unused invalid token".to_string());
|
|
assert!(validate_bgfilter_internal_token_for_startup(&controller_config).is_ok());
|
|
|
|
let missing_config = AppConfig::default();
|
|
assert!(validate_bgfilter_internal_token_for_startup(&missing_config).is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn auth_store_startup_restore_is_limited_to_http_roles() {
|
|
assert!(should_restore_auth_store_for_startup(ProcessRole::Api));
|
|
assert!(should_restore_auth_store_for_startup(ProcessRole::All));
|
|
assert!(!should_restore_auth_store_for_startup(
|
|
ProcessRole::BgfilterWorker
|
|
));
|
|
assert!(!should_restore_auth_store_for_startup(
|
|
ProcessRole::ExternalGenerationWorker
|
|
));
|
|
assert!(!should_restore_auth_store_for_startup(
|
|
ProcessRole::ExternalGenerationController
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn editor_generation_pricing_initialization_is_limited_to_http_roles() {
|
|
assert!(should_initialize_editor_generation_pricing_for_startup(
|
|
ProcessRole::Api
|
|
));
|
|
assert!(should_initialize_editor_generation_pricing_for_startup(
|
|
ProcessRole::All
|
|
));
|
|
assert!(!should_initialize_editor_generation_pricing_for_startup(
|
|
ProcessRole::BgfilterWorker
|
|
));
|
|
assert!(!should_initialize_editor_generation_pricing_for_startup(
|
|
ProcessRole::ExternalGenerationWorker
|
|
));
|
|
assert!(!should_initialize_editor_generation_pricing_for_startup(
|
|
ProcessRole::ExternalGenerationController
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn profile_recharge_expiration_listener_is_limited_to_http_roles() {
|
|
assert!(should_start_profile_recharge_expiration_listener(
|
|
ProcessRole::Api
|
|
));
|
|
assert!(should_start_profile_recharge_expiration_listener(
|
|
ProcessRole::All
|
|
));
|
|
assert!(!should_start_profile_recharge_expiration_listener(
|
|
ProcessRole::BgfilterWorker
|
|
));
|
|
assert!(!should_start_profile_recharge_expiration_listener(
|
|
ProcessRole::ExternalGenerationWorker
|
|
));
|
|
assert!(!should_start_profile_recharge_expiration_listener(
|
|
ProcessRole::ExternalGenerationController
|
|
));
|
|
}
|
|
}
|