feat(server,dashboard): 引入多工作流数据隔离、安全中间件与前端 ESM 模块化重构
- server: 实现按 workflow_name 的多工作流数据隔离与旧数据库平滑迁移机制 - server: 新增 API Key 认证(auth)、限流中间件(rate_limit)与运维备份接口(admin) - server: 统一 AppError 错误处理体系,重构调度器 scheduler 支持工作流级重置与抢占 - node: 节点 ID 缺失时自动生成随机 UUID,原生支持 `docker compose --scale node=N` 动态扩容 - dashboard: 前端模块化重构(state/api/components),升级 CSS 变量设计系统与 Toast 通知 - docker/docs: 更新 /healthz 健康检查、部署脚本 IP 配置及数据库设计文档
This commit is contained in:
+112
-28
@@ -135,19 +135,56 @@ impl GridConfig {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[derive(Clone, Serialize, Deserialize)]
|
||||
pub struct ServerConfig {
|
||||
pub bind_addr: String,
|
||||
pub db_path: String,
|
||||
pub queue_db_path: String,
|
||||
pub results_dir: String,
|
||||
/// 数据库备份目录(每日自动备份落盘位置)。默认 "data/backups",可经 DCTS_BACKUP_DIR 覆盖。
|
||||
pub backup_dir: String,
|
||||
pub grid_config: String,
|
||||
pub stale_sec: u64,
|
||||
#[serde(default = "default_node_stale_sec")]
|
||||
pub node_stale_sec: u64,
|
||||
pub mq_type: String, // "sqlite" or "rabbitmq"
|
||||
pub rabbitmq_url: Option<String>,
|
||||
/// Admin 凭据(Dashboard 登录用)。优先 DCTS_ADMIN_TOKEN,回退旧变量 DCTS_AUTH_TOKEN。
|
||||
#[serde(default)]
|
||||
pub admin_token: Option<String>,
|
||||
/// 兼容字段:保留以判断「是否启用鉴权」与旧中间件逻辑。取 admin_token 的值。
|
||||
#[serde(default)]
|
||||
pub auth_token: Option<String>,
|
||||
/// 应急开关:DCTS_AUTH_DISABLE=1 时跳过全部鉴权(仅本地调试,默认关闭)。
|
||||
#[serde(default)]
|
||||
pub auth_disabled: bool,
|
||||
}
|
||||
|
||||
// 手写 Debug:token 类字段脱敏为 ***REDACTED***,防止日志/错误链泄露明文凭据。
|
||||
impl std::fmt::Debug for ServerConfig {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("ServerConfig")
|
||||
.field("bind_addr", &self.bind_addr)
|
||||
.field("db_path", &self.db_path)
|
||||
.field("queue_db_path", &self.queue_db_path)
|
||||
.field("results_dir", &self.results_dir)
|
||||
.field("backup_dir", &self.backup_dir)
|
||||
.field("grid_config", &self.grid_config)
|
||||
.field("stale_sec", &self.stale_sec)
|
||||
.field("node_stale_sec", &self.node_stale_sec)
|
||||
.field("mq_type", &self.mq_type)
|
||||
.field("rabbitmq_url", &self.rabbitmq_url)
|
||||
.field(
|
||||
"admin_token",
|
||||
&self.admin_token.as_ref().map(|_| "***REDACTED***"),
|
||||
)
|
||||
.field(
|
||||
"auth_token",
|
||||
&self.auth_token.as_ref().map(|_| "***REDACTED***"),
|
||||
)
|
||||
.field("auth_disabled", &self.auth_disabled)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
fn default_node_stale_sec() -> u64 {
|
||||
@@ -160,12 +197,13 @@ impl Default for ServerConfig {
|
||||
.or_else(|_| std::env::var("CNO_PORT"))
|
||||
.or_else(|_| std::env::var("PORT"))
|
||||
.unwrap_or_else(|_| "8090".to_string());
|
||||
let db_path = std::env::var("DCTS_DB_PATH")
|
||||
.unwrap_or_else(|_| "data/dcts.db".to_string());
|
||||
let db_path = std::env::var("DCTS_DB_PATH").unwrap_or_else(|_| "data/dcts.db".to_string());
|
||||
let queue_db_path = std::env::var("DCTS_QUEUE_DB_PATH")
|
||||
.unwrap_or_else(|_| "data/dcts_queue.db".to_string());
|
||||
let results_dir = std::env::var("DCTS_RESULTS_DIR")
|
||||
.unwrap_or_else(|_| "data/results".to_string());
|
||||
let results_dir =
|
||||
std::env::var("DCTS_RESULTS_DIR").unwrap_or_else(|_| "data/results".to_string());
|
||||
let backup_dir =
|
||||
std::env::var("DCTS_BACKUP_DIR").unwrap_or_else(|_| "data/backups".to_string());
|
||||
let grid_config = std::env::var("DCTS_GRID_CONFIG")
|
||||
.unwrap_or_else(|_| "workflows/sdB_cno.yaml".to_string());
|
||||
// 默认设置为 7800 秒,比计算任务默认超时(7200 秒)高 600 秒缓冲,避免两边的超时检测同时触发冲突
|
||||
@@ -177,26 +215,63 @@ impl Default for ServerConfig {
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<u64>().ok())
|
||||
.unwrap_or(60);
|
||||
let mq_type = std::env::var("DCTS_MQ_TYPE")
|
||||
.unwrap_or_else(|_| "sqlite".to_string());
|
||||
let mq_type = std::env::var("DCTS_MQ_TYPE").unwrap_or_else(|_| "sqlite".to_string());
|
||||
let rabbitmq_url = std::env::var("DCTS_RABBITMQ_URL").ok();
|
||||
let auth_token = std::env::var("DCTS_AUTH_TOKEN").ok();
|
||||
|
||||
// ── 鉴权凭据解析 ──
|
||||
let legacy_token = std::env::var("DCTS_AUTH_TOKEN")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty());
|
||||
let admin_token = std::env::var("DCTS_ADMIN_TOKEN")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
.or_else(|| legacy_token.clone());
|
||||
if legacy_token.is_some()
|
||||
&& (std::env::var("DCTS_ADMIN_TOKEN").is_err()
|
||||
|| std::env::var("DCTS_ENROLLMENT_TOKEN").is_err())
|
||||
{
|
||||
tracing::warn!(
|
||||
"检测到旧的 DCTS_AUTH_TOKEN,已自动用作 admin/enrollment 凭据。\
|
||||
建议迁移到 DCTS_ADMIN_TOKEN(管理)与 DCTS_ENROLLMENT_TOKEN(节点注册)"
|
||||
);
|
||||
}
|
||||
|
||||
let auth_disabled = std::env::var("DCTS_AUTH_DISABLE")
|
||||
.map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
|
||||
.unwrap_or(false);
|
||||
if auth_disabled {
|
||||
tracing::warn!(
|
||||
"⚠️ DCTS_AUTH_DISABLE=1 已生效:全部鉴权被跳过,仅供本地调试,切勿用于生产!"
|
||||
);
|
||||
}
|
||||
|
||||
// auth_token 兼容字段:用于 main.rs 判断「是否启用鉴权中间件」。
|
||||
// 启用条件 = 显式配置了 admin 或 enrollment 凭据,且未应急关闭。
|
||||
let auth_token = if auth_disabled {
|
||||
None
|
||||
} else {
|
||||
admin_token.clone()
|
||||
};
|
||||
|
||||
Self {
|
||||
bind_addr: format!("0.0.0.0:{}", port),
|
||||
db_path,
|
||||
queue_db_path,
|
||||
results_dir,
|
||||
backup_dir,
|
||||
grid_config,
|
||||
stale_sec,
|
||||
node_stale_sec,
|
||||
mq_type,
|
||||
rabbitmq_url,
|
||||
admin_token,
|
||||
auth_token,
|
||||
auth_disabled,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[derive(Clone, Serialize, Deserialize)]
|
||||
pub struct NodeConfig {
|
||||
pub node_id: String,
|
||||
pub server_url: String,
|
||||
@@ -204,7 +279,19 @@ pub struct NodeConfig {
|
||||
pub runtime_dir: String,
|
||||
pub work_dir: String,
|
||||
pub heartbeat_sec: u64,
|
||||
pub auth_token: Option<String>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for NodeConfig {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("NodeConfig")
|
||||
.field("node_id", &self.node_id)
|
||||
.field("server_url", &self.server_url)
|
||||
.field("max_slots", &self.max_slots)
|
||||
.field("runtime_dir", &self.runtime_dir)
|
||||
.field("work_dir", &self.work_dir)
|
||||
.field("heartbeat_sec", &self.heartbeat_sec)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for NodeConfig {
|
||||
@@ -215,22 +302,27 @@ impl Default for NodeConfig {
|
||||
.unwrap_or_else(|_| "http://127.0.0.1:8090".to_string());
|
||||
let node_id = std::env::var("DCTS_NODE_ID")
|
||||
.or_else(|_| std::env::var("NODE_ID"))
|
||||
.and_then(|v| if v.trim().is_empty() { Err(std::env::VarError::NotPresent) } else { Ok(v) })
|
||||
.and_then(|v| {
|
||||
if v.trim().is_empty() {
|
||||
Err(std::env::VarError::NotPresent)
|
||||
} else {
|
||||
Ok(v)
|
||||
}
|
||||
})
|
||||
.unwrap_or_else(|_| format!("node-{}", uuid::Uuid::new_v4().simple()));
|
||||
let max_slots = std::env::var("DCTS_MAX_SLOTS")
|
||||
.or_else(|_| std::env::var("MAX_SLOTS"))
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<usize>().ok())
|
||||
.unwrap_or(4);
|
||||
let runtime_dir = std::env::var("DCTS_RUNTIME_DIR")
|
||||
.unwrap_or_else(|_| "data/runtime".to_string());
|
||||
let work_dir = std::env::var("DCTS_WORK_DIR")
|
||||
.unwrap_or_else(|_| "data/work".to_string());
|
||||
let runtime_dir =
|
||||
std::env::var("DCTS_RUNTIME_DIR").unwrap_or_else(|_| "data/runtime".to_string());
|
||||
let work_dir = std::env::var("DCTS_WORK_DIR").unwrap_or_else(|_| "data/work".to_string());
|
||||
let heartbeat_sec = std::env::var("DCTS_HEARTBEAT_SEC")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<u64>().ok())
|
||||
.unwrap_or(15);
|
||||
let auth_token = std::env::var("DCTS_AUTH_TOKEN").ok();
|
||||
|
||||
Self {
|
||||
node_id,
|
||||
server_url,
|
||||
@@ -238,7 +330,6 @@ impl Default for NodeConfig {
|
||||
runtime_dir,
|
||||
work_dir,
|
||||
heartbeat_sec,
|
||||
auth_token,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -250,21 +341,14 @@ mod tests {
|
||||
#[test]
|
||||
fn test_load_real_grid_configs() {
|
||||
let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../..");
|
||||
|
||||
|
||||
let sdb_path = root.join("workflows/sdB_cno.yaml");
|
||||
if sdb_path.exists() {
|
||||
let cfg = GridConfig::load_from_file(&sdb_path).expect("解析 workflows/sdB_cno.yaml 发生失败");
|
||||
let cfg = GridConfig::load_from_file(&sdb_path)
|
||||
.expect("解析 workflows/sdB_cno.yaml 发生失败");
|
||||
assert_eq!(cfg.nworkers, 16);
|
||||
assert_eq!(cfg.niter, Some(100));
|
||||
}
|
||||
|
||||
let dense_path = root.join("config_dense.yaml");
|
||||
if dense_path.exists() {
|
||||
let cfg = GridConfig::load_from_file(&dense_path).expect("解析 config_dense.yaml 发生失败");
|
||||
assert_eq!(cfg.template.as_deref(), Some("templates/cno_atmos.5.tpl"));
|
||||
assert_eq!(cfg.linelist.as_deref(), Some("data/gfVIS99.dat"));
|
||||
assert!(cfg.seed_step_fallback);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -80,15 +80,12 @@ pub fn check_fort9(path: &Path, chmax: f64) -> ConvCheckResult {
|
||||
}
|
||||
|
||||
// Safely find depth with maximum absolute change without unwrap panic on NaN
|
||||
let worst = match cur_rows
|
||||
.iter()
|
||||
.max_by(|a, b| {
|
||||
a.maximum
|
||||
.abs()
|
||||
.partial_cmp(&b.maximum.abs())
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
{
|
||||
let worst = match cur_rows.iter().max_by(|a, b| {
|
||||
a.maximum
|
||||
.abs()
|
||||
.partial_cmp(&b.maximum.abs())
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
}) {
|
||||
Some(row) => row,
|
||||
None => {
|
||||
return ConvCheckResult {
|
||||
@@ -98,7 +95,9 @@ pub fn check_fort9(path: &Path, chmax: f64) -> ConvCheckResult {
|
||||
last_iter,
|
||||
n_depths: 0,
|
||||
chmax,
|
||||
error: Some("No valid iteration rows found when calculating maximum change".to_string()),
|
||||
error: Some(
|
||||
"No valid iteration rows found when calculating maximum change".to_string(),
|
||||
),
|
||||
};
|
||||
}
|
||||
};
|
||||
@@ -113,15 +112,22 @@ pub fn check_fort9(path: &Path, chmax: f64) -> ConvCheckResult {
|
||||
last_iter,
|
||||
n_depths: cur_rows.len(),
|
||||
chmax,
|
||||
error: if is_valid_num { None } else { Some("Convergence value is NaN or Inf".to_string()) },
|
||||
error: if is_valid_num {
|
||||
None
|
||||
} else {
|
||||
Some("Convergence value is NaN or Inf".to_string())
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Checks if an atmosphere file (.7) contains NaN lines (>10% NaN lines = invalid) using exact word boundary
|
||||
///
|
||||
/// 文件缺失时返回 `false`(语义:不存在 NaN 内容)。这与“含 NaN 导致无效”是不同语义;
|
||||
/// 调用方需先自行确认文件存在性,不应将“缺失”与“含 NaN”混为一谈。
|
||||
pub fn atmosphere_has_nan(path: &Path) -> bool {
|
||||
let file = match File::open(path) {
|
||||
Ok(f) => f,
|
||||
Err(_) => return true,
|
||||
Err(_) => return false,
|
||||
};
|
||||
let reader = BufReader::new(file);
|
||||
let mut total_lines = 0;
|
||||
@@ -142,7 +148,6 @@ pub fn atmosphere_has_nan(path: &Path) -> bool {
|
||||
(nan_lines as f64) > (total_lines as f64 * 0.1)
|
||||
}
|
||||
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -162,5 +167,9 @@ mod tests {
|
||||
let banana_file_path = dir.path().join("banana.7");
|
||||
std::fs::write(&banana_file_path, "banana 2 3\nbanana 5 6\n7 8 9\n").unwrap();
|
||||
assert!(!atmosphere_has_nan(&banana_file_path));
|
||||
|
||||
// Missing file returns false (absence != contains NaN)
|
||||
let missing_path = dir.path().join("missing.7");
|
||||
assert!(!atmosphere_has_nan(&missing_path));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -54,9 +54,13 @@ pub async fn ensure_runtime(
|
||||
write_if_changed(&synspec_exe, SYNSPEC_BIN, true)?;
|
||||
}
|
||||
|
||||
|
||||
// 2. Fetch baseline equation of state partition function tables if missing locally
|
||||
let common_files = &["irwin_bc.dat", "irwin_orig.dat", "tsuji.molec_bc2", "tsuji.molec_orig"];
|
||||
let common_files = &[
|
||||
"irwin_bc.dat",
|
||||
"irwin_orig.dat",
|
||||
"tsuji.molec_bc2",
|
||||
"tsuji.molec_orig",
|
||||
];
|
||||
ensure_specific_data_files(&data_dir, server_url, client, common_files).await?;
|
||||
|
||||
// 3. Check gfVIS99.dat
|
||||
@@ -69,11 +73,15 @@ pub async fn ensure_runtime(
|
||||
fs::write(&linelist, &bytes)?;
|
||||
info!("成功下载并保存主谱线库 gfVIS99.dat");
|
||||
} else {
|
||||
anyhow::bail!("从服务端下载主谱线库 gfVIS99.dat 失败,HTTP 状态码: {}", resp.status());
|
||||
anyhow::bail!(
|
||||
"从服务端下载主谱线库 gfVIS99.dat 失败,HTTP 状态码: {}",
|
||||
resp.status()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let abs_runtime_dir = fs::canonicalize(runtime_dir).unwrap_or_else(|_| runtime_dir.to_path_buf());
|
||||
let abs_runtime_dir =
|
||||
fs::canonicalize(runtime_dir).unwrap_or_else(|_| runtime_dir.to_path_buf());
|
||||
let tlusty_exe = abs_runtime_dir.join("tlusty_static");
|
||||
let synspec_exe = abs_runtime_dir.join("synspec_static");
|
||||
let data_dir = abs_runtime_dir.join("data");
|
||||
@@ -103,17 +111,28 @@ pub async fn ensure_specific_data_files(
|
||||
let local_file = data_dir.join(filename);
|
||||
if !local_file.exists() {
|
||||
let file_url = format!("{}/api/data/file/{}", server_url, filename);
|
||||
info!("本地缺失数据文件 {},开始从服务端拉取: {}", filename, file_url);
|
||||
info!(
|
||||
"本地缺失数据文件 {},开始从服务端拉取: {}",
|
||||
filename, file_url
|
||||
);
|
||||
|
||||
let resp = client.get(&file_url).send().await?;
|
||||
if resp.status().is_success() {
|
||||
let bytes = resp.bytes().await?;
|
||||
let tmp_file = data_dir.join(format!("{}.{}.tmp", filename, uuid::Uuid::new_v4().simple()));
|
||||
let tmp_file = data_dir.join(format!(
|
||||
"{}.{}.tmp",
|
||||
filename,
|
||||
uuid::Uuid::new_v4().simple()
|
||||
));
|
||||
tokio::fs::write(&tmp_file, &bytes).await?;
|
||||
tokio::fs::rename(&tmp_file, &local_file).await?;
|
||||
info!("成功保存数据文件: {}", filename);
|
||||
} else {
|
||||
anyhow::bail!("服务端返回 HTTP {} 错误,数据文件: {}", resp.status(), filename);
|
||||
anyhow::bail!(
|
||||
"服务端返回 HTTP {} 错误,数据文件: {}",
|
||||
resp.status(),
|
||||
filename
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,10 +7,16 @@ pub fn generate_fort55_content(cfg: &SynspecConfig) -> String {
|
||||
let line3 = " 0 0 0 0 0";
|
||||
let line4 = " 1 1 0 0 0";
|
||||
let line5 = " 0 0 0";
|
||||
let line6 = format!(" {:.1} {:.1} 10 0 {} {}", cfg.wstart, cfg.wend, cfg.rel_cutoff, cfg.abs_cutoff);
|
||||
let line6 = format!(
|
||||
" {:.1} {:.1} 10 0 {} {}",
|
||||
cfg.wstart, cfg.wend, cfg.rel_cutoff, cfg.abs_cutoff
|
||||
);
|
||||
let line7 = " 0 0";
|
||||
|
||||
format!("{}\n{}\n{}\n{}\n{}\n{}\n{}\n", line1, line2, line3, line4, line5, line6, line7)
|
||||
format!(
|
||||
"{}\n{}\n{}\n{}\n{}\n{}\n{}\n",
|
||||
line1, line2, line3, line4, line5, line6, line7
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
+161
-28
@@ -9,40 +9,172 @@ struct IonDef {
|
||||
}
|
||||
|
||||
const IONS_H: &[IonDef] = &[
|
||||
IonDef { iat: 1, iz: 0, nlevs: 9, typion: " H 1", filei: "data/h1.dat" },
|
||||
IonDef { iat: 1, iz: 1, nlevs: 1, typion: " H 2", filei: " " },
|
||||
IonDef {
|
||||
iat: 1,
|
||||
iz: 0,
|
||||
nlevs: 9,
|
||||
typion: " H 1",
|
||||
filei: "data/h1.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 1,
|
||||
iz: 1,
|
||||
nlevs: 1,
|
||||
typion: " H 2",
|
||||
filei: " ",
|
||||
},
|
||||
];
|
||||
|
||||
const IONS_HE: &[IonDef] = &[
|
||||
IonDef { iat: 2, iz: 0, nlevs: 14, typion: "He 1", filei: "data/he1.dat" },
|
||||
IonDef { iat: 2, iz: 1, nlevs: 14, typion: "He 2", filei: "data/he2.dat" },
|
||||
IonDef { iat: 2, iz: 2, nlevs: 1, typion: "He 3", filei: " " },
|
||||
IonDef {
|
||||
iat: 2,
|
||||
iz: 0,
|
||||
nlevs: 14,
|
||||
typion: "He 1",
|
||||
filei: "data/he1.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 2,
|
||||
iz: 1,
|
||||
nlevs: 14,
|
||||
typion: "He 2",
|
||||
filei: "data/he2.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 2,
|
||||
iz: 2,
|
||||
nlevs: 1,
|
||||
typion: "He 3",
|
||||
filei: " ",
|
||||
},
|
||||
];
|
||||
|
||||
const IONS_C: &[IonDef] = &[
|
||||
IonDef { iat: 6, iz: 0, nlevs: 40, typion: " C 1", filei: "data/c1.dat" },
|
||||
IonDef { iat: 6, iz: 1, nlevs: 22, typion: " C 2", filei: "data/c2.dat" },
|
||||
IonDef { iat: 6, iz: 2, nlevs: 46, typion: " C 3", filei: "data/c3_34+12lev.dat" },
|
||||
IonDef { iat: 6, iz: 3, nlevs: 25, typion: " C 4", filei: "data/c4.dat" },
|
||||
IonDef { iat: 6, iz: 4, nlevs: 1, typion: " C 5", filei: " " },
|
||||
IonDef {
|
||||
iat: 6,
|
||||
iz: 0,
|
||||
nlevs: 40,
|
||||
typion: " C 1",
|
||||
filei: "data/c1.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 6,
|
||||
iz: 1,
|
||||
nlevs: 22,
|
||||
typion: " C 2",
|
||||
filei: "data/c2.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 6,
|
||||
iz: 2,
|
||||
nlevs: 46,
|
||||
typion: " C 3",
|
||||
filei: "data/c3_34+12lev.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 6,
|
||||
iz: 3,
|
||||
nlevs: 25,
|
||||
typion: " C 4",
|
||||
filei: "data/c4.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 6,
|
||||
iz: 4,
|
||||
nlevs: 1,
|
||||
typion: " C 5",
|
||||
filei: " ",
|
||||
},
|
||||
];
|
||||
|
||||
const IONS_N: &[IonDef] = &[
|
||||
IonDef { iat: 7, iz: 0, nlevs: 34, typion: " N 1", filei: "data/n1.dat" },
|
||||
IonDef { iat: 7, iz: 1, nlevs: 42, typion: " N 2", filei: "data/n2_32+10lev.dat" },
|
||||
IonDef { iat: 7, iz: 2, nlevs: 32, typion: " N 3", filei: "data/n3.dat" },
|
||||
IonDef { iat: 7, iz: 3, nlevs: 48, typion: " N 4", filei: "data/n4_34+14lev.dat" },
|
||||
IonDef { iat: 7, iz: 4, nlevs: 16, typion: " N 5", filei: "data/n5.dat" },
|
||||
IonDef { iat: 7, iz: 5, nlevs: 1, typion: " N 6", filei: " " },
|
||||
IonDef {
|
||||
iat: 7,
|
||||
iz: 0,
|
||||
nlevs: 34,
|
||||
typion: " N 1",
|
||||
filei: "data/n1.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 7,
|
||||
iz: 1,
|
||||
nlevs: 42,
|
||||
typion: " N 2",
|
||||
filei: "data/n2_32+10lev.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 7,
|
||||
iz: 2,
|
||||
nlevs: 32,
|
||||
typion: " N 3",
|
||||
filei: "data/n3.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 7,
|
||||
iz: 3,
|
||||
nlevs: 48,
|
||||
typion: " N 4",
|
||||
filei: "data/n4_34+14lev.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 7,
|
||||
iz: 4,
|
||||
nlevs: 16,
|
||||
typion: " N 5",
|
||||
filei: "data/n5.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 7,
|
||||
iz: 5,
|
||||
nlevs: 1,
|
||||
typion: " N 6",
|
||||
filei: " ",
|
||||
},
|
||||
];
|
||||
|
||||
const IONS_O: &[IonDef] = &[
|
||||
IonDef { iat: 8, iz: 0, nlevs: 33, typion: " O 1", filei: "data/o1_23+10lev.dat" },
|
||||
IonDef { iat: 8, iz: 1, nlevs: 48, typion: " O 2", filei: "data/o2_36+12lev.dat" },
|
||||
IonDef { iat: 8, iz: 2, nlevs: 41, typion: " O 3", filei: "data/o3_28+13lev.dat" },
|
||||
IonDef { iat: 8, iz: 3, nlevs: 39, typion: " O 4", filei: "data/o4.dat" },
|
||||
IonDef { iat: 8, iz: 4, nlevs: 6, typion: " O 5", filei: "data/o5.dat" },
|
||||
IonDef { iat: 8, iz: 5, nlevs: 1, typion: " O 6", filei: " " },
|
||||
IonDef {
|
||||
iat: 8,
|
||||
iz: 0,
|
||||
nlevs: 33,
|
||||
typion: " O 1",
|
||||
filei: "data/o1_23+10lev.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 8,
|
||||
iz: 1,
|
||||
nlevs: 48,
|
||||
typion: " O 2",
|
||||
filei: "data/o2_36+12lev.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 8,
|
||||
iz: 2,
|
||||
nlevs: 41,
|
||||
typion: " O 3",
|
||||
filei: "data/o3_28+13lev.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 8,
|
||||
iz: 3,
|
||||
nlevs: 39,
|
||||
typion: " O 4",
|
||||
filei: "data/o4.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 8,
|
||||
iz: 4,
|
||||
nlevs: 6,
|
||||
typion: " O 5",
|
||||
filei: "data/o5.dat",
|
||||
},
|
||||
IonDef {
|
||||
iat: 8,
|
||||
iz: 5,
|
||||
nlevs: 1,
|
||||
typion: " O 6",
|
||||
filei: " ",
|
||||
},
|
||||
];
|
||||
|
||||
fn fmt_abn(logx: f64) -> String {
|
||||
@@ -64,11 +196,11 @@ pub fn make_input5(
|
||||
|
||||
// Atoms block
|
||||
let mut atom_rows: Vec<(i32, String)> = vec![
|
||||
(2, "0.".to_string()), // 1 H
|
||||
(2, fmt_abn(params.loghe)), // 2 He
|
||||
(0, "0.".to_string()), // 3 Li
|
||||
(0, "0.".to_string()), // 4 Be
|
||||
(0, "0.".to_string()), // 5 B
|
||||
(2, "0.".to_string()), // 1 H
|
||||
(2, fmt_abn(params.loghe)), // 2 He
|
||||
(0, "0.".to_string()), // 3 Li
|
||||
(0, "0.".to_string()), // 4 Be
|
||||
(0, "0.".to_string()), // 5 B
|
||||
];
|
||||
|
||||
if has_c {
|
||||
@@ -81,7 +213,8 @@ pub fn make_input5(
|
||||
atom_rows.push((2, fmt_abn(params.logo))); // 8 O
|
||||
}
|
||||
|
||||
let natoms = 5 + (if has_c { 1 } else { 0 }) + (if has_n { 1 } else { 0 }) + (if has_o { 1 } else { 0 });
|
||||
let natoms =
|
||||
5 + (if has_c { 1 } else { 0 }) + (if has_n { 1 } else { 0 }) + (if has_o { 1 } else { 0 });
|
||||
|
||||
let mut atoms_block = format!(" {}\n* mode abn modpf\n", natoms);
|
||||
for (mode, abn) in &atom_rows {
|
||||
|
||||
@@ -8,4 +8,3 @@ pub mod models;
|
||||
pub mod nst_writer;
|
||||
pub mod runner;
|
||||
pub mod seed_finder;
|
||||
|
||||
|
||||
@@ -17,7 +17,6 @@ impl FormatTime for LocalTimeFormatter {
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Initializes high-performance, non-blocking structured logging for DCTS applications
|
||||
pub fn init_logging(app_name: &str, default_filter: &str) -> Result<Vec<WorkerGuard>> {
|
||||
let mut guards = Vec::new();
|
||||
@@ -29,7 +28,8 @@ pub fn init_logging(app_name: &str, default_filter: &str) -> Result<Vec<WorkerGu
|
||||
let log_outputs = env::var("LOG_OUTPUTS").unwrap_or_else(|_| "stdout,file".to_string());
|
||||
let log_dir = env::var("LOG_DIR").unwrap_or_else(|_| "data/logs".to_string());
|
||||
|
||||
let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(&log_level));
|
||||
let env_filter =
|
||||
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(&log_level));
|
||||
let is_json = log_format.to_lowercase() == "json";
|
||||
|
||||
let mut layers: Vec<Box<dyn Layer<tracing_subscriber::Registry> + Send + Sync>> = Vec::new();
|
||||
@@ -39,7 +39,9 @@ pub fn init_logging(app_name: &str, default_filter: &str) -> Result<Vec<WorkerGu
|
||||
let (non_blocking, guard) = tracing_appender::non_blocking(std::io::stdout());
|
||||
guards.push(guard);
|
||||
|
||||
let fmt_layer = fmt::layer().with_timer(LocalTimeFormatter).with_writer(non_blocking);
|
||||
let fmt_layer = fmt::layer()
|
||||
.with_timer(LocalTimeFormatter)
|
||||
.with_writer(non_blocking);
|
||||
if is_json {
|
||||
layers.push(fmt_layer.json().with_ansi(false).boxed());
|
||||
} else {
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
/// 6D grid point parameter specification
|
||||
@@ -84,7 +84,6 @@ impl std::fmt::Display for GridPointStatus {
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
impl From<&str> for GridPointStatus {
|
||||
fn from(s: &str) -> Self {
|
||||
match s {
|
||||
@@ -106,6 +105,10 @@ pub struct TaskSpec {
|
||||
pub task_type: TaskType,
|
||||
pub seed_point_name: Option<String>,
|
||||
pub timeout_sec: u64,
|
||||
/// 所属工作流名称,用于按工作流隔离队列清理(stop_workflow 只清当前工作流的任务)。
|
||||
/// 旧数据反序列化时缺省为 None。
|
||||
#[serde(default)]
|
||||
pub workflow_name: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
@@ -251,10 +254,12 @@ mod tests {
|
||||
assert_eq!(GridPointStatus::Failed.to_string(), "failed");
|
||||
|
||||
assert_eq!(GridPointStatus::from("queued"), GridPointStatus::Queued);
|
||||
assert_eq!(GridPointStatus::from("converged"), GridPointStatus::Converged);
|
||||
assert_eq!(
|
||||
GridPointStatus::from("converged"),
|
||||
GridPointStatus::Converged
|
||||
);
|
||||
assert_eq!(GridPointStatus::from("done"), GridPointStatus::Converged);
|
||||
assert_eq!(GridPointStatus::from("failed"), GridPointStatus::Failed);
|
||||
assert_eq!(GridPointStatus::from("unknown"), GridPointStatus::Pending);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -63,4 +63,3 @@ mod tests {
|
||||
assert!(content.contains("IELCOR=-1"));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+72
-16
@@ -6,11 +6,11 @@ use crate::gen_input5::make_input5;
|
||||
use crate::models::{GridPointParams, ModelSummary, StageSummary, TaskType};
|
||||
use crate::nst_writer::generate_nst_content;
|
||||
use anyhow::Result;
|
||||
use tokio::fs::File;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Stdio;
|
||||
use tokio::process::Command as AsyncCommand;
|
||||
use std::time::Instant;
|
||||
use tokio::fs::File;
|
||||
use tokio::process::Command as AsyncCommand;
|
||||
use tracing::{info, warn};
|
||||
|
||||
pub fn default_cold_chain() -> Vec<StageConfig> {
|
||||
@@ -130,7 +130,15 @@ impl<'a> ExecutionRunner<'a> {
|
||||
seed_atmos: Option<&Path>,
|
||||
synspec_cfg: Option<&SynspecConfig>,
|
||||
) -> Result<ModelSummary> {
|
||||
self.run_model_with_timeout(params, task_type, custom_chain, seed_atmos, synspec_cfg, 7200).await
|
||||
self.run_model_with_timeout(
|
||||
params,
|
||||
task_type,
|
||||
custom_chain,
|
||||
seed_atmos,
|
||||
synspec_cfg,
|
||||
7200,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn run_model_with_timeout(
|
||||
@@ -157,7 +165,9 @@ impl<'a> ExecutionRunner<'a> {
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let abs_data_dir = tokio::fs::canonicalize(&self.runtime.data_dir).await.unwrap_or_else(|_| self.runtime.data_dir.clone());
|
||||
let abs_data_dir = tokio::fs::canonicalize(&self.runtime.data_dir)
|
||||
.await
|
||||
.unwrap_or_else(|_| self.runtime.data_dir.clone());
|
||||
if let Err(e) = std::os::unix::fs::symlink(&abs_data_dir, &link_data) {
|
||||
warn!("构建 data 数据集软链时发生提示性告警: {}", e);
|
||||
}
|
||||
@@ -171,7 +181,10 @@ impl<'a> ExecutionRunner<'a> {
|
||||
if let Some(seed_path) = seed_atmos {
|
||||
if seed_path.is_file() {
|
||||
if let Err(e) = tokio::fs::copy(seed_path, &fort8).await {
|
||||
warn!("向工作沙盒引导填载首期收敛模型种子 fort.8 发生复制错误: {}", e);
|
||||
warn!(
|
||||
"向工作沙盒引导填载首期收敛模型种子 fort.8 发生复制错误: {}",
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -219,15 +232,24 @@ impl<'a> ExecutionRunner<'a> {
|
||||
} else if let Some(ref s_path) = current_seed {
|
||||
if s_path.is_file() {
|
||||
if let Err(e) = tokio::fs::copy(s_path, &fort8).await {
|
||||
warn!("阶段 {} 重载候选近邻推算种子模型期间发生文件复制异常: {}", stage_def.label, e);
|
||||
warn!(
|
||||
"阶段 {} 重载候选近邻推算种子模型期间发生文件复制异常: {}",
|
||||
stage_def.label, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Run tlusty.exe
|
||||
let fin = File::open(&input5_path).await?.into_std().await;
|
||||
let fout = File::create(model_dir.join(format!("{}.6", name))).await?.into_std().await;
|
||||
let ferr = File::create(model_dir.join(format!("{}.err", name))).await?.into_std().await;
|
||||
let fout = File::create(model_dir.join(format!("{}.6", name)))
|
||||
.await?
|
||||
.into_std()
|
||||
.await;
|
||||
let ferr = File::create(model_dir.join(format!("{}.err", name)))
|
||||
.await?
|
||||
.into_std()
|
||||
.await;
|
||||
|
||||
let child = AsyncCommand::new(&self.runtime.tlusty_exe)
|
||||
.current_dir(&model_dir)
|
||||
@@ -246,7 +268,6 @@ impl<'a> ExecutionRunner<'a> {
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
let fort9 = model_dir.join("fort.9");
|
||||
let fort7 = model_dir.join("fort.7");
|
||||
|
||||
@@ -293,7 +314,10 @@ impl<'a> ExecutionRunner<'a> {
|
||||
stage_summaries.push(stage_summary);
|
||||
|
||||
if !final_converged && stage_def.require_converged {
|
||||
warn!("阶段 {} 要求收敛但未达标,中止后续收敛链阶段", stage_def.label);
|
||||
warn!(
|
||||
"阶段 {} 要求收敛但未达标,中止后续收敛链阶段",
|
||||
stage_def.label
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -344,14 +368,19 @@ impl<'a> ExecutionRunner<'a> {
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let abs_linelist = tokio::fs::canonicalize(&self.runtime.linelist).await.unwrap_or_else(|_| self.runtime.linelist.clone());
|
||||
let abs_linelist = tokio::fs::canonicalize(&self.runtime.linelist)
|
||||
.await
|
||||
.unwrap_or_else(|_| self.runtime.linelist.clone());
|
||||
let _ = std::os::unix::fs::symlink(&abs_linelist, &fort19_path);
|
||||
}
|
||||
|
||||
let input5_path = model_dir.join(format!("{}.5", name));
|
||||
if input5_path.is_file() {
|
||||
let fin = File::open(&input5_path).await?.into_std().await;
|
||||
let fout = File::create(model_dir.join(format!("{}.log", name))).await?.into_std().await;
|
||||
let fout = File::create(model_dir.join(format!("{}.log", name)))
|
||||
.await?
|
||||
.into_std()
|
||||
.await;
|
||||
|
||||
let child = AsyncCommand::new(&self.runtime.synspec_exe)
|
||||
.current_dir(&model_dir)
|
||||
@@ -361,7 +390,8 @@ impl<'a> ExecutionRunner<'a> {
|
||||
.kill_on_drop(true)
|
||||
.spawn()?;
|
||||
|
||||
let status_res = run_child_async_with_timeout(child, timeout_sec).await;
|
||||
let synspec_timeout_sec = 600_u64.min(timeout_sec);
|
||||
let status_res = run_child_async_with_timeout(child, synspec_timeout_sec).await;
|
||||
let rc = match status_res {
|
||||
Ok(st) => st.code().unwrap_or(-1),
|
||||
Err(e) => {
|
||||
@@ -374,13 +404,25 @@ impl<'a> ExecutionRunner<'a> {
|
||||
|
||||
// Copy/move outputs: fort.7 (Synspec spectrum) -> .spec, fort.17 -> .cont, fort.12 -> .iden
|
||||
if model_dir.join("fort.7").is_file() {
|
||||
let _ = tokio::fs::rename(model_dir.join("fort.7"), model_dir.join(format!("{}.spec", name))).await;
|
||||
let _ = tokio::fs::rename(
|
||||
model_dir.join("fort.7"),
|
||||
model_dir.join(format!("{}.spec", name)),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
if model_dir.join("fort.17").is_file() {
|
||||
let _ = tokio::fs::copy(model_dir.join("fort.17"), model_dir.join(format!("{}.cont", name))).await;
|
||||
let _ = tokio::fs::copy(
|
||||
model_dir.join("fort.17"),
|
||||
model_dir.join(format!("{}.cont", name)),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
if model_dir.join("fort.12").is_file() {
|
||||
let _ = tokio::fs::copy(model_dir.join("fort.12"), model_dir.join(format!("{}.iden", name))).await;
|
||||
let _ = tokio::fs::copy(
|
||||
model_dir.join("fort.12"),
|
||||
model_dir.join(format!("{}.iden", name)),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
@@ -415,3 +457,17 @@ impl<'a> ExecutionRunner<'a> {
|
||||
Ok(summary)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[test]
|
||||
fn test_synspec_timeout_calculation() {
|
||||
let long_tlusty_timeout: u64 = 7200;
|
||||
let synspec_timeout = 600_u64.min(long_tlusty_timeout);
|
||||
assert_eq!(synspec_timeout, 600);
|
||||
|
||||
let short_tlusty_timeout: u64 = 300;
|
||||
let synspec_timeout_short = 600_u64.min(short_tlusty_timeout);
|
||||
assert_eq!(synspec_timeout_short, 300);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -18,7 +18,11 @@ pub fn calculate_seed_distance(cand: &GridPointParams, target: &GridPointParams)
|
||||
+ (cand.logn - target.logn).abs()
|
||||
+ (cand.logo - target.logo).abs();
|
||||
|
||||
if d_teff < 1.0 && d_logg < 0.01 && d_loghe < 0.01 {
|
||||
// exact family 判定:Teff/logg/logHe 视为“同物理族”,仅 CNO 丰度不同。
|
||||
// Teff 容忍度取半步 5000K:实际网格 Teff 档位通常为整数千(20000/30000/.../60000),
|
||||
// 半步既能覆盖 config_dense 等 10000K 步长的相邻档互作种子,
|
||||
// 又避免跨过大 Teff 间距导致 sdB 高温模型用低温种子而不收敛(sdB_cno 步长 40000K 仍不命中 exact)。
|
||||
if d_teff < 5000.0 && d_logg < 0.01 && d_loghe < 0.01 {
|
||||
(true, d_cno)
|
||||
} else {
|
||||
// 距离公式物理意义与标定阐释:
|
||||
@@ -32,6 +36,3 @@ pub fn calculate_seed_distance(cand: &GridPointParams, target: &GridPointParams)
|
||||
(false, global_d)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,2 +1 @@
|
||||
pub mod sqlite_queue;
|
||||
|
||||
|
||||
+219
-27
@@ -10,11 +10,32 @@ struct SqliteCustomizer;
|
||||
|
||||
impl r2d2::CustomizeConnection<rusqlite::Connection, rusqlite::Error> for SqliteCustomizer {
|
||||
fn on_acquire(&self, conn: &mut rusqlite::Connection) -> Result<(), rusqlite::Error> {
|
||||
conn.pragma_update(None, "busy_timeout", 5000)?;
|
||||
// 与主库一致:高并发 claim/report 下给 SQLITE_BUSY 足够重试窗口。
|
||||
conn.pragma_update(None, "busy_timeout", 15000)?;
|
||||
conn.pragma_update(None, "wal_autocheckpoint", 1000)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 SQLite db 文件及其 WAL/SHM 侧车文件权限收紧为 0600(仅 owner 读写)。
|
||||
/// 与主库 dcts.db 的口径一致,作为纵深防御(队列库不含 token,但含任务 payload)。
|
||||
#[cfg(unix)]
|
||||
fn restrict_db_file_perms(db_path: &str) {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let candidates = [
|
||||
std::path::PathBuf::from(db_path),
|
||||
std::path::PathBuf::from(format!("{}-wal", db_path)),
|
||||
std::path::PathBuf::from(format!("{}-shm", db_path)),
|
||||
];
|
||||
for p in candidates {
|
||||
if let Ok(meta) = std::fs::metadata(&p) {
|
||||
let mut perms = meta.permissions();
|
||||
perms.set_mode(0o600);
|
||||
let _ = std::fs::set_permissions(&p, perms);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct SqliteTaskQueue {
|
||||
pool: Pool<SqliteConnectionManager>,
|
||||
@@ -29,11 +50,15 @@ impl SqliteTaskQueue {
|
||||
}
|
||||
let manager = SqliteConnectionManager::file(&db_path_owned);
|
||||
let pool = Pool::builder()
|
||||
.max_size(4)
|
||||
.max_size(8)
|
||||
.connection_customizer(Box::new(SqliteCustomizer))
|
||||
.build(manager)
|
||||
.context("Failed to build SQLite queue connection pool")?;
|
||||
|
||||
// 收紧队列 db 文件权限为 0600(与主库口径一致,纵深防御)
|
||||
#[cfg(unix)]
|
||||
restrict_db_file_perms(&db_path_owned);
|
||||
|
||||
let conn = pool.get()?;
|
||||
let _: String = conn.pragma_update_and_check(None, "journal_mode", "WAL", |r| r.get(0))?;
|
||||
conn.execute(
|
||||
@@ -42,14 +67,39 @@ impl SqliteTaskQueue {
|
||||
payload TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL,
|
||||
claimed_at DATETIME
|
||||
claimed_at DATETIME,
|
||||
workflow_name TEXT,
|
||||
claimed_by_node_id TEXT
|
||||
)",
|
||||
[],
|
||||
)?;
|
||||
// 兼容旧库:若 task_queue 表已存在但缺少 workflow_name / claimed_by_node_id 列,则补列。
|
||||
// SQLite 的 ALTER TABLE ADD COLUMN 是在线操作,旧数据该列默认 NULL。
|
||||
// PRAGMA table_info 检测列是否存在以实现幂等 migration。
|
||||
let has_col = |conn: &rusqlite::Connection, col: &str| -> rusqlite::Result<bool> {
|
||||
let mut stmt = conn.prepare("PRAGMA table_info(task_queue)")?;
|
||||
let rows = stmt.query_map([], |r| r.get::<_, String>(1))?;
|
||||
for r in rows {
|
||||
if r.map(|name| name == col).unwrap_or(false) {
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
Ok(false)
|
||||
};
|
||||
if !has_col(&conn, "workflow_name")? {
|
||||
conn.execute("ALTER TABLE task_queue ADD COLUMN workflow_name TEXT", [])?;
|
||||
}
|
||||
if !has_col(&conn, "claimed_by_node_id")? {
|
||||
conn.execute("ALTER TABLE task_queue ADD COLUMN claimed_by_node_id TEXT", [])?;
|
||||
}
|
||||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_task_queue_status_created ON task_queue(status, created_at)",
|
||||
[],
|
||||
)?;
|
||||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_task_queue_workflow ON task_queue(workflow_name)",
|
||||
[],
|
||||
)?;
|
||||
Ok(pool)
|
||||
})
|
||||
.await??;
|
||||
@@ -61,14 +111,15 @@ impl SqliteTaskQueue {
|
||||
pub async fn push_task(&self, task: &TaskSpec) -> Result<()> {
|
||||
let payload = serde_json::to_string(task)?;
|
||||
let task_id_str = task.task_id.to_string();
|
||||
let workflow_name = task.workflow_name.clone();
|
||||
let pool = self.pool.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO task_queue (task_id, payload, status, created_at)
|
||||
VALUES (?1, ?2, 'pending', datetime('now'))",
|
||||
params![task_id_str, payload],
|
||||
"INSERT OR REPLACE INTO task_queue (task_id, payload, status, created_at, workflow_name)
|
||||
VALUES (?1, ?2, 'pending', datetime('now'), ?3)",
|
||||
params![task_id_str, payload, workflow_name],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
@@ -77,8 +128,9 @@ impl SqliteTaskQueue {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn pop_task(&self) -> Result<Option<TaskSpec>> {
|
||||
pub async fn pop_task(&self, claimant_node_id: &str) -> Result<Option<TaskSpec>> {
|
||||
let pool = self.pool.clone();
|
||||
let claimant = claimant_node_id.to_string();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<Option<TaskSpec>> {
|
||||
let mut attempts = 0;
|
||||
@@ -109,9 +161,11 @@ impl SqliteTaskQueue {
|
||||
|
||||
let task: TaskSpec = serde_json::from_str(&payload)?;
|
||||
|
||||
// 记录任务归属:claim 时写入领用方 node_id,供 report 阶段校验,
|
||||
// 杜绝「节点 A 领用、节点 B 上报」的跨节点伪造结果投毒。
|
||||
tx.execute(
|
||||
"UPDATE task_queue SET status = 'claimed', claimed_at = datetime('now') WHERE task_id = ?1",
|
||||
params![task_id],
|
||||
"UPDATE task_queue SET status = 'claimed', claimed_at = datetime('now'), claimed_by_node_id = ?2 WHERE task_id = ?1",
|
||||
params![task_id, claimant],
|
||||
)?;
|
||||
|
||||
tx.commit()?;
|
||||
@@ -139,8 +193,13 @@ impl SqliteTaskQueue {
|
||||
let id_owned = task_id.to_string();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
conn.execute("DELETE FROM task_queue WHERE task_id = ?1", params![id_owned])?;
|
||||
let conn = pool
|
||||
.get()
|
||||
.map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
conn.execute(
|
||||
"DELETE FROM task_queue WHERE task_id = ?1",
|
||||
params![id_owned],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
@@ -148,13 +207,60 @@ impl SqliteTaskQueue {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn requeue_stale_tasks(&self, stale_sec: u64) -> Result<Vec<String>> {
|
||||
/// 校验指定 task 是否由指定 node 领用(claimed 态且 claimed_by_node_id 匹配)。
|
||||
///
|
||||
/// 用于 report_task 阶段防止跨节点伪造结果:只有真正领用该 task 的 node 才能上报结果。
|
||||
/// 返回 (point_name, workflow_name):匹配时附带二者供 report 进一步校验「上报的点与领用的
|
||||
/// task 一致」并把 workflow_name 传给 record_task_report 以定向更新对应工作流的 grid_points
|
||||
/// (多工作流分区:避免按 name 全局更新误改其他工作流同名点)。
|
||||
pub async fn verify_task_claim(
|
||||
&self,
|
||||
task_id: &str,
|
||||
claimant_node_id: &str,
|
||||
) -> Result<Option<(String, Option<String>)>> {
|
||||
let pool = self.pool.clone();
|
||||
let task_id = task_id.to_string();
|
||||
let claimant = claimant_node_id.to_string();
|
||||
|
||||
let res =
|
||||
tokio::task::spawn_blocking(move || -> Result<Option<(String, Option<String>)>> {
|
||||
let conn = pool
|
||||
.get()
|
||||
.map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
// 仅 claimed 态(尚未被 report 清理)且归属匹配才算有效领用
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT payload FROM task_queue
|
||||
WHERE task_id = ?1 AND claimed_by_node_id = ?2 AND status = 'claimed' LIMIT 1",
|
||||
)?;
|
||||
let row = stmt.query_row(params![task_id, claimant], |r| r.get::<_, String>(0));
|
||||
match row {
|
||||
Ok(payload) => {
|
||||
// 解析 payload 取出 point_name + workflow_name,供调用方校验与定向更新
|
||||
let task: TaskSpec = serde_json::from_str(&payload)?;
|
||||
Ok(Some((task.point_name, task.workflow_name)))
|
||||
}
|
||||
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
})
|
||||
.await??;
|
||||
Ok(res)
|
||||
}
|
||||
|
||||
/// 重投超时 claimed 任务回 pending,返回每个被重投任务的 (point_name, workflow_name)。
|
||||
///
|
||||
/// 返回 workflow_name 供调用方(main.rs 后台循环)按工作流分组调用
|
||||
/// reset_specific_grid_points_to_pending,避免跨工作流误改同名点(多工作流分区)。
|
||||
pub async fn requeue_stale_tasks(
|
||||
&self,
|
||||
stale_sec: u64,
|
||||
) -> Result<Vec<(String, Option<String>)>> {
|
||||
let pool = self.pool.clone();
|
||||
|
||||
let names = tokio::task::spawn_blocking(move || -> Result<Vec<String>> {
|
||||
let entries = tokio::task::spawn_blocking(move || -> Result<Vec<(String, Option<String>)>> {
|
||||
let mut conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
let tx = conn.transaction()?;
|
||||
let mut point_names = Vec::new();
|
||||
let mut entries = Vec::new();
|
||||
|
||||
{
|
||||
// 改写为单一原子更新带 RETURNING 返回语句,消弭 TOCTOU (Time-Of-Check-To-Time-Of-Use) 竞态问题
|
||||
@@ -165,32 +271,56 @@ impl SqliteTaskQueue {
|
||||
)?;
|
||||
let rows = stmt.query_map(params![stale_sec as i64], |row| row.get::<_, String>(0))?;
|
||||
for r in rows {
|
||||
if let Ok(payload) = r {
|
||||
if let Ok(task) = serde_json::from_str::<TaskSpec>(&payload) {
|
||||
point_names.push(task.point_name);
|
||||
}
|
||||
let payload = match r {
|
||||
Ok(p) => p,
|
||||
Err(_) => continue,
|
||||
};
|
||||
if let Ok(task) = serde_json::from_str::<TaskSpec>(&payload) {
|
||||
entries.push((task.point_name, task.workflow_name));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tx.commit()?;
|
||||
Ok(point_names)
|
||||
Ok(entries)
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(names)
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
pub async fn clear_queue(&self) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
let conn = pool
|
||||
.get()
|
||||
.map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
conn.execute("DELETE FROM task_queue", [])?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 仅清理指定工作流的排队任务。
|
||||
///
|
||||
/// 用于 stop_workflow 按工作流隔离清理,避免在多工作流场景下误清其他工作流的任务。
|
||||
pub async fn clear_queue_by_workflow(&self, workflow_name: &str) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let wf_owned = workflow_name.to_string();
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool
|
||||
.get()
|
||||
.map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
conn.execute(
|
||||
"DELETE FROM task_queue WHERE workflow_name = ?1",
|
||||
params![wf_owned],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -203,9 +333,11 @@ mod tests {
|
||||
async fn test_sqlite_task_queue_operations() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let db_path = temp_dir.path().join("test_queue.db");
|
||||
let queue = SqliteTaskQueue::new(&db_path.to_string_lossy()).await.unwrap();
|
||||
let queue = SqliteTaskQueue::new(&db_path.to_string_lossy())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(queue.pop_task().await.unwrap().is_none());
|
||||
assert!(queue.pop_task("test-node").await.unwrap().is_none());
|
||||
|
||||
let task_id = Uuid::new_v4();
|
||||
let task = TaskSpec {
|
||||
@@ -222,24 +354,84 @@ mod tests {
|
||||
task_type: TaskType::ColdRun,
|
||||
seed_point_name: None,
|
||||
timeout_sec: 3600,
|
||||
workflow_name: Some("test_wf".to_string()),
|
||||
};
|
||||
queue.push_task(&task).await.unwrap();
|
||||
|
||||
let popped = queue.pop_task().await.unwrap();
|
||||
let popped = queue.pop_task("test-node").await.unwrap();
|
||||
assert!(popped.is_some());
|
||||
let popped_task = popped.unwrap();
|
||||
assert_eq!(popped_task.task_id, task_id);
|
||||
assert_eq!(popped_task.point_name, task.point_name);
|
||||
|
||||
assert!(queue.pop_task().await.unwrap().is_none());
|
||||
assert!(queue.pop_task("test-node").await.unwrap().is_none());
|
||||
|
||||
let requeued = queue.requeue_stale_tasks(0).await.unwrap();
|
||||
assert_eq!(requeued.len(), 1);
|
||||
|
||||
let popped2 = queue.pop_task().await.unwrap();
|
||||
let popped2 = queue.pop_task("test-node").await.unwrap();
|
||||
assert!(popped2.is_some());
|
||||
|
||||
queue.remove_task(&task_id.to_string()).await.unwrap();
|
||||
assert!(queue.pop_task().await.unwrap().is_none());
|
||||
assert!(queue.pop_task("test-node").await.unwrap().is_none());
|
||||
}
|
||||
|
||||
/// 任务归属校验:领用方 node 匹配才放行,其他 node 校验失败(防跨节点伪造结果)。
|
||||
#[tokio::test]
|
||||
async fn test_verify_task_claim_ownership() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let db_path = temp_dir.path().join("claim_test.db");
|
||||
let queue = SqliteTaskQueue::new(&db_path.to_string_lossy())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let task_id = Uuid::new_v4();
|
||||
let params = GridPointParams {
|
||||
teff: 35000.0,
|
||||
logg: 5.5,
|
||||
loghe: -1.0,
|
||||
logc: -2.0,
|
||||
logn: -2.0,
|
||||
logo: -2.0,
|
||||
};
|
||||
let task = TaskSpec {
|
||||
task_id,
|
||||
point_name: params.model_name(),
|
||||
params: params.clone(),
|
||||
task_type: TaskType::ColdRun,
|
||||
seed_point_name: None,
|
||||
timeout_sec: 60,
|
||||
workflow_name: None,
|
||||
};
|
||||
queue.push_task(&task).await.unwrap();
|
||||
|
||||
// node-A 领用
|
||||
let popped = queue.pop_task("node-A").await.unwrap();
|
||||
assert!(popped.is_some());
|
||||
|
||||
// node-A 校验:匹配,返回绑定的 (point_name, workflow_name)
|
||||
let claim = queue
|
||||
.verify_task_claim(&task_id.to_string(), "node-A")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
claim.map(|(p, _)| p).as_deref(),
|
||||
Some(params.model_name().as_str())
|
||||
);
|
||||
|
||||
// node-B 校验:非领用方,返回 None
|
||||
let claim_b = queue
|
||||
.verify_task_claim(&task_id.to_string(), "node-B")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(claim_b.is_none());
|
||||
|
||||
// 任务被清理(remove)后,任何 node 校验都失败
|
||||
queue.remove_task(&task_id.to_string()).await.unwrap();
|
||||
let claim_after = queue
|
||||
.verify_task_claim(&task_id.to_string(), "node-A")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(claim_after.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
+128
-13
@@ -13,17 +13,35 @@ pub async fn execute_task(
|
||||
work_dir: &Path,
|
||||
task: &TaskSpec,
|
||||
) -> Result<(ModelSummary, Option<Vec<u8>>)> {
|
||||
info!("开始执行计算任务 {} (网格点: {})", task.task_id, task.point_name);
|
||||
info!(
|
||||
"开始执行计算任务 {} (网格点: {})",
|
||||
task.task_id, task.point_name
|
||||
);
|
||||
|
||||
// 1. Pull ONLY missing atom model data files needed for this task
|
||||
let required_atom_files = &[
|
||||
"h1.dat", "he1.dat", "he2.dat",
|
||||
"c1.dat", "c2.dat", "c3_34+12lev.dat", "c4.dat",
|
||||
"n1.dat", "n2_32+10lev.dat", "n3.dat", "n4_34+14lev.dat", "n5.dat",
|
||||
"o1_23+10lev.dat", "o2_36+12lev.dat", "o3_28+13lev.dat", "o4.dat", "o5.dat",
|
||||
"h1.dat",
|
||||
"he1.dat",
|
||||
"he2.dat",
|
||||
"c1.dat",
|
||||
"c2.dat",
|
||||
"c3_34+12lev.dat",
|
||||
"c4.dat",
|
||||
"n1.dat",
|
||||
"n2_32+10lev.dat",
|
||||
"n3.dat",
|
||||
"n4_34+14lev.dat",
|
||||
"n5.dat",
|
||||
"o1_23+10lev.dat",
|
||||
"o2_36+12lev.dat",
|
||||
"o3_28+13lev.dat",
|
||||
"o4.dat",
|
||||
"o5.dat",
|
||||
];
|
||||
|
||||
if let Err(e) = ensure_specific_data_files(&runtime.data_dir, server_url, client, required_atom_files).await {
|
||||
if let Err(e) =
|
||||
ensure_specific_data_files(&runtime.data_dir, server_url, client, required_atom_files).await
|
||||
{
|
||||
warn!("拉取缺失原子数据文件失败: {}", e);
|
||||
}
|
||||
|
||||
@@ -40,7 +58,15 @@ pub async fn execute_task(
|
||||
if let Ok(bytes) = resp.bytes().await {
|
||||
let temp_seed_dir = work_dir.join(".seed_cache");
|
||||
tokio::fs::create_dir_all(&temp_seed_dir).await?;
|
||||
let tmp_path = temp_seed_dir.join(format!("{}.{}.tmp", seed_name, uuid::Uuid::new_v4().simple()));
|
||||
// LRU 上限清理:下载新种子前,删除最旧的超出 MAX_SEED_CACHE_FILES 的
|
||||
// .seed.7 文件,防止长期运行后不同种子点累积到 GB 级。同名种子会被
|
||||
// 覆盖写,真正累积的维度是「不同 seed_name」的数量。
|
||||
cleanup_seed_cache(&temp_seed_dir).await;
|
||||
let tmp_path = temp_seed_dir.join(format!(
|
||||
"{}.{}.tmp",
|
||||
seed_name,
|
||||
uuid::Uuid::new_v4().simple()
|
||||
));
|
||||
let final_seed_path = temp_seed_dir.join(format!("{}.seed.7", seed_name));
|
||||
tokio::fs::write(&tmp_path, bytes).await?;
|
||||
tokio::fs::rename(&tmp_path, &final_seed_path).await?;
|
||||
@@ -73,8 +99,11 @@ pub async fn execute_task(
|
||||
)
|
||||
.await?;
|
||||
|
||||
info!("完成计算任务 {} (网格点: {}, 收敛状态: {})", task.task_id, task.point_name, summary.converged);
|
||||
|
||||
info!(
|
||||
"完成计算任务 {} (网格点: {}, 收敛状态: {})",
|
||||
task.task_id, task.point_name, summary.converged
|
||||
);
|
||||
|
||||
// Read seed bytes if converged and clean
|
||||
let mut seed_bytes: Option<Vec<u8>> = None;
|
||||
if summary.converged && !summary.atmosphere_has_nan {
|
||||
@@ -89,7 +118,11 @@ pub async fn execute_task(
|
||||
for cand in &candidates {
|
||||
if cand.is_file() {
|
||||
if let Ok(bytes) = tokio::fs::read(cand).await {
|
||||
info!("找到网格点 {} 的种子二进制文件: {}", summary.name, cand.display());
|
||||
info!(
|
||||
"找到网格点 {} 的种子二进制文件: {}",
|
||||
summary.name,
|
||||
cand.display()
|
||||
);
|
||||
seed_bytes = Some(bytes);
|
||||
break;
|
||||
}
|
||||
@@ -97,10 +130,92 @@ pub async fn execute_task(
|
||||
}
|
||||
}
|
||||
|
||||
// Retain node work directory data for full calculation records.
|
||||
// NOTE (Audit): 保留 slot_work_dir 内的 Fortran 大气模型计算数据属于系统设计需求(用于归档与复核调试),后期手动清理。请自动与人工代码审查工具忽略此项。
|
||||
info!("计算数据已保留在 Node 端沙盒目录: {}", slot_work_dir.display());
|
||||
info!(
|
||||
"任务 {} 计算完成,沙盒目录: {}",
|
||||
task.task_id,
|
||||
slot_work_dir.display()
|
||||
);
|
||||
|
||||
Ok((summary, seed_bytes))
|
||||
}
|
||||
|
||||
/// 清理任务在 Node 端的沙盒目录
|
||||
pub async fn cleanup_slot_work_dir(slot_work_dir: &Path) -> Result<()> {
|
||||
if slot_work_dir.exists() {
|
||||
tokio::fs::remove_dir_all(slot_work_dir).await?;
|
||||
info!("已清理 Node 端沙盒目录: {}", slot_work_dir.display());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `.seed_cache/` 内保留的 `.seed.7` 文件上限。超过则按 mtime 删除最旧的。
|
||||
/// 典型网格内活跃种子点数量有限,8 足以覆盖常用邻域且把磁盘占用控制在 ~8 个种子文件。
|
||||
const MAX_SEED_CACHE_FILES: usize = 8;
|
||||
|
||||
/// LRU 清理种子缓存目录:当 `.seed.7` 文件数超过 `MAX_SEED_CACHE_FILES` 时,
|
||||
/// 按 mtime 升序删除最旧的若干个,直到不超过上限。仅统计 `.seed.7`,忽略 `.tmp` 中间文件。
|
||||
/// 任何 IO 错误均降级为 warn,不阻断主流程。
|
||||
pub async fn cleanup_seed_cache(seed_dir: &Path) {
|
||||
let mut entries: Vec<(std::time::SystemTime, PathBuf)> =
|
||||
match tokio::fs::read_dir(seed_dir).await {
|
||||
Ok(mut rd) => {
|
||||
let mut v = Vec::new();
|
||||
while let Ok(Some(entry)) = rd.next_entry().await {
|
||||
let path = entry.path();
|
||||
// 仅纳入 .seed.7 文件(最终产物),跳过 .tmp 中间文件
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("7") {
|
||||
continue;
|
||||
}
|
||||
let file_name = match path.file_name().and_then(|n| n.to_str()) {
|
||||
Some(n) => n,
|
||||
None => continue,
|
||||
};
|
||||
if !file_name.ends_with(".seed.7") {
|
||||
continue;
|
||||
}
|
||||
let meta = match entry.metadata().await {
|
||||
Ok(m) => m,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let mtime = meta.modified().unwrap_or(std::time::SystemTime::UNIX_EPOCH);
|
||||
v.push((mtime, path));
|
||||
}
|
||||
v
|
||||
}
|
||||
Err(_) => return,
|
||||
};
|
||||
|
||||
if entries.len() <= MAX_SEED_CACHE_FILES {
|
||||
return;
|
||||
}
|
||||
|
||||
// 按 mtime 升序(最旧在前),删除超出上限的最旧文件
|
||||
entries.sort_by_key(|(mtime, _)| *mtime);
|
||||
let to_remove = entries.len().saturating_sub(MAX_SEED_CACHE_FILES);
|
||||
for (_, path) in entries.into_iter().take(to_remove) {
|
||||
if let Err(e) = tokio::fs::remove_file(&path).await {
|
||||
warn!("清理种子缓存文件 {} 失败: {}", path.display(), e);
|
||||
} else {
|
||||
info!("LRU 清理种子缓存文件: {}", path.display());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cleanup_slot_work_dir() {
|
||||
let temp_dir =
|
||||
std::env::temp_dir().join(format!("test_slot_work_dir_{}", uuid::Uuid::new_v4()));
|
||||
tokio::fs::create_dir_all(&temp_dir).await.unwrap();
|
||||
tokio::fs::write(temp_dir.join("dummy.txt"), "content")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(temp_dir.exists());
|
||||
cleanup_slot_work_dir(&temp_dir).await.unwrap();
|
||||
assert!(!temp_dir.exists());
|
||||
}
|
||||
}
|
||||
|
||||
+73
-10
@@ -8,7 +8,7 @@ use common::embedded::ensure_runtime;
|
||||
use common::logging::init_logging;
|
||||
use reqwest::Client;
|
||||
use std::path::Path;
|
||||
use tracing::info;
|
||||
use tracing::{info, warn};
|
||||
use worker::NodeWorker;
|
||||
|
||||
#[tokio::main]
|
||||
@@ -19,17 +19,36 @@ async fn main() -> Result<()> {
|
||||
info!("启动 DCTS 计算节点 (Distributed Computing TLUSTY/SYNSPEC Worker Node)...");
|
||||
|
||||
let node_cfg = NodeConfig::default();
|
||||
|
||||
let runtime_dir = Path::new(&node_cfg.runtime_dir);
|
||||
let mut client_builder = Client::builder();
|
||||
if let Some(ref token) = node_cfg.auth_token {
|
||||
let mut headers = reqwest::header::HeaderMap::new();
|
||||
if let Ok(val) = reqwest::header::HeaderValue::from_str(&format!("Bearer {}", token)) {
|
||||
headers.insert(reqwest::header::AUTHORIZATION, val);
|
||||
|
||||
// ── Node 凭据获取 ──
|
||||
// 1. 优先读取本地持久化的 node 专属 token(`.node_token`,权限 600)。
|
||||
// 2. 若不存在,调 /node/register 提交注册申请并轮询等待 Dashboard 管理员审批授权。
|
||||
let token_path = runtime_dir.join(".node_token");
|
||||
let node_token = match read_node_token(&token_path) {
|
||||
Some(t) => {
|
||||
info!("已加载本地持久化的 node 专属 token");
|
||||
t
|
||||
}
|
||||
client_builder = client_builder.default_headers(headers);
|
||||
}
|
||||
let client = client_builder.build().unwrap_or_else(|_| Client::new());
|
||||
None => {
|
||||
info!("本地未发现 node token,准备向服务端提交注册申请并等待管理员审批...");
|
||||
let public_client = build_client_with_token(None);
|
||||
let issued = NodeWorker::register_and_fetch_token(
|
||||
&public_client,
|
||||
&node_cfg.server_url,
|
||||
&node_cfg.node_id,
|
||||
)
|
||||
.await
|
||||
.context("向服务端提交申请或获取专属 token 失败")?;
|
||||
|
||||
write_node_token(&token_path, &issued)?;
|
||||
info!("已持久化获批的专属 node token 到 {}", token_path.display());
|
||||
issued
|
||||
}
|
||||
};
|
||||
|
||||
// 用 node 专属 token 构造后续所有请求的 client
|
||||
let client = build_client_with_token(Some(&node_token));
|
||||
|
||||
info!("检查本地运行时二进制与基础数据文件,必要时从服务端拉取...");
|
||||
let runtime = ensure_runtime(runtime_dir, &node_cfg.server_url, &client)
|
||||
@@ -41,3 +60,47 @@ async fn main() -> Result<()> {
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 读取本地持久化的 node token;文件须存在且非空。
|
||||
fn read_node_token(path: &Path) -> Option<String> {
|
||||
let content = std::fs::read_to_string(path).ok()?;
|
||||
let t = content.trim().to_string();
|
||||
if t.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(t)
|
||||
}
|
||||
}
|
||||
|
||||
/// 持久化 node token 到本地文件,并设权限 600(仅 owner 可读写)。
|
||||
fn write_node_token(path: &Path, token: &str) -> Result<()> {
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)
|
||||
.with_context(|| format!("创建 token 目录失败: {}", parent.display()))?;
|
||||
}
|
||||
std::fs::write(path, token)
|
||||
.with_context(|| format!("写入 token 文件失败: {}", path.display()))?;
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let mut perms = std::fs::metadata(path)?.permissions();
|
||||
perms.set_mode(0o600);
|
||||
std::fs::set_permissions(path, perms)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 构造带 Authorization: Bearer 头的 reqwest client。
|
||||
fn build_client_with_token(token: Option<&str>) -> Client {
|
||||
let mut builder = Client::builder();
|
||||
if let Some(t) = token {
|
||||
let mut headers = reqwest::header::HeaderMap::new();
|
||||
if let Ok(val) = reqwest::header::HeaderValue::from_str(&format!("Bearer {}", t)) {
|
||||
headers.insert(reqwest::header::AUTHORIZATION, val);
|
||||
} else {
|
||||
warn!("node token 含非法 HTTP 头字符,已忽略鉴权头");
|
||||
}
|
||||
builder = builder.default_headers(headers);
|
||||
}
|
||||
builder.build().unwrap_or_else(|_| Client::new())
|
||||
}
|
||||
|
||||
+39
-30
@@ -4,7 +4,6 @@ use reqwest::multipart::{Form, Part};
|
||||
use reqwest::Client;
|
||||
use tracing::{info, warn};
|
||||
|
||||
|
||||
pub async fn report_result(
|
||||
client: &Client,
|
||||
server_url: &str,
|
||||
@@ -14,32 +13,33 @@ pub async fn report_result(
|
||||
) -> Result<()> {
|
||||
let report_url = format!("{}/api/task/report", server_url);
|
||||
|
||||
let (status, converged, max_relc, atmo_has_nan, elapsed_sec, err_msg, summary_json, seed_bytes) = match exec_res {
|
||||
Ok((s, s_bytes)) => (
|
||||
if s.converged {
|
||||
TaskStatus::Completed
|
||||
} else {
|
||||
TaskStatus::Failed
|
||||
},
|
||||
s.converged,
|
||||
s.final_max_relc,
|
||||
s.atmosphere_has_nan,
|
||||
s.elapsed_sec,
|
||||
s.note.clone(),
|
||||
serde_json::to_string(&s).unwrap_or_default(),
|
||||
s_bytes,
|
||||
),
|
||||
Err(e) => (
|
||||
TaskStatus::Failed,
|
||||
false,
|
||||
None,
|
||||
false,
|
||||
0.0,
|
||||
Some(e.clone()),
|
||||
serde_json::json!({"error": e}).to_string(),
|
||||
None,
|
||||
),
|
||||
};
|
||||
let (status, converged, max_relc, atmo_has_nan, elapsed_sec, err_msg, summary_json, seed_bytes) =
|
||||
match exec_res {
|
||||
Ok((s, s_bytes)) => (
|
||||
if s.converged {
|
||||
TaskStatus::Completed
|
||||
} else {
|
||||
TaskStatus::Failed
|
||||
},
|
||||
s.converged,
|
||||
s.final_max_relc,
|
||||
s.atmosphere_has_nan,
|
||||
s.elapsed_sec,
|
||||
s.note.clone(),
|
||||
serde_json::to_string(&s).unwrap_or_default(),
|
||||
s_bytes,
|
||||
),
|
||||
Err(e) => (
|
||||
TaskStatus::Failed,
|
||||
false,
|
||||
None,
|
||||
false,
|
||||
0.0,
|
||||
Some(e.clone()),
|
||||
serde_json::json!({"error": e}).to_string(),
|
||||
None,
|
||||
),
|
||||
};
|
||||
|
||||
let report = TaskReport {
|
||||
task_id: task.task_id,
|
||||
@@ -83,9 +83,19 @@ pub async fn report_result(
|
||||
return Ok(());
|
||||
}
|
||||
Ok(resp) => {
|
||||
let status = resp.status();
|
||||
// 401/403 表明 node token 已失效/被吊销(非临时故障),重试无意义且会丢结果。
|
||||
// 立即 bail 并打 error,与 claim_task 侧口径统一,提示运维介入。
|
||||
if status.as_u16() == 401 || status.as_u16() == 403 {
|
||||
tracing::error!(
|
||||
"上报任务 {} 被服务端拒绝 (HTTP {}):node token 可能已失效或被吊销,请检查并清理 .node_token 文件后重启节点以重新向服务端发起注册审批,停止重试",
|
||||
task.task_id, status
|
||||
);
|
||||
anyhow::bail!("node token 失效或被吊销 (HTTP {}),结果未上报", status);
|
||||
}
|
||||
warn!(
|
||||
"向服务端上报任务 {} 结果失败 (尝试 {}/{}): HTTP {}",
|
||||
task.task_id, attempt, max_attempts, resp.status()
|
||||
task.task_id, attempt, max_attempts, status
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
@@ -97,7 +107,7 @@ pub async fn report_result(
|
||||
}
|
||||
|
||||
if attempt < max_attempts {
|
||||
let backoff_secs = (1 << (attempt - 1)).min(60);
|
||||
let backoff_secs = (1u64 << (attempt - 1).min(6)).min(60);
|
||||
let backoff = std::time::Duration::from_secs(backoff_secs);
|
||||
tokio::time::sleep(backoff).await;
|
||||
}
|
||||
@@ -109,4 +119,3 @@ pub async fn report_result(
|
||||
task.task_id
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
+162
-21
@@ -29,8 +29,96 @@ impl NodeWorker {
|
||||
}
|
||||
}
|
||||
|
||||
/// 仅注册并领取专属 token(供 main.rs 在本地无 token 时调用)。
|
||||
/// 支持免凭据申请注册并轮询等待管理员在 Web Dashboard 上点击同意。
|
||||
pub async fn register_and_fetch_token(
|
||||
client: &Client,
|
||||
server_url: &str,
|
||||
node_id: &str,
|
||||
) -> Result<String> {
|
||||
info!(
|
||||
"正在向服务端 {} 提交计算节点 {} 的注册申请...",
|
||||
server_url, node_id
|
||||
);
|
||||
|
||||
let req = NodeRegisterRequest {
|
||||
node_id: node_id.to_string(),
|
||||
host_name: gethostname::gethostname().to_string_lossy().to_string(),
|
||||
max_slots: 0,
|
||||
};
|
||||
|
||||
let resp = client
|
||||
.post(format!("{}/api/node/register", server_url))
|
||||
.json(&req)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
anyhow::bail!("向服务端提交注册申请失败,HTTP 状态码: {}", resp.status());
|
||||
}
|
||||
|
||||
let json: Value = resp.json().await?;
|
||||
let status = json.get("status").and_then(|v| v.as_str()).unwrap_or("");
|
||||
|
||||
if status == "approved" {
|
||||
if let Some(t) = json.get("node_token").and_then(|v| v.as_str()) {
|
||||
return Ok(t.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
info!(
|
||||
"⏳ 节点 {} 的注册申请已提交!等待管理员在管理 Dashboard 上点击【同意接入】...",
|
||||
node_id
|
||||
);
|
||||
|
||||
// 轮询等待管理员在 Dashboard 上的 Approve
|
||||
loop {
|
||||
sleep(Duration::from_secs(5)).await;
|
||||
|
||||
let check_req = serde_json::json!({ "node_id": node_id });
|
||||
let resp = match client
|
||||
.post(format!("{}/api/node/check_status", server_url))
|
||||
.json(&check_req)
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
warn!("轮询节点审批状态网络异常: {}", e);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if !resp.status().is_success() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let body: Value = match resp.json().await {
|
||||
Ok(b) => b,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
let check_status = body.get("status").and_then(|v| v.as_str()).unwrap_or("");
|
||||
if check_status == "approved" {
|
||||
if let Some(token) = body.get("node_token").and_then(|v| v.as_str()) {
|
||||
info!(
|
||||
"🎉 节点 {} 已成功获取管理员授权!专属访问 Token 接收完成。",
|
||||
node_id
|
||||
);
|
||||
return Ok(token.to_string());
|
||||
}
|
||||
} else if check_status == "rejected" {
|
||||
anyhow::bail!("节点 {} 的注册申请已被管理员拒绝或清理", node_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 正式注册(带真实 slot 数),供 run() 启动时刷新节点信息用。
|
||||
pub async fn register(&self) -> Result<()> {
|
||||
info!("正在向服务端 {} 注册计算节点 {}...", self.config.server_url, self.config.node_id);
|
||||
info!(
|
||||
"正在向服务端 {} 刷新节点 {} 注册信息...",
|
||||
self.config.server_url, self.config.node_id
|
||||
);
|
||||
|
||||
let req = NodeRegisterRequest {
|
||||
node_id: self.config.node_id.clone(),
|
||||
@@ -54,7 +142,10 @@ impl NodeWorker {
|
||||
|
||||
pub async fn run(&self) -> Result<()> {
|
||||
self.register().await?;
|
||||
info!("计算节点已激活,最大并行 Slot 槽位数: {}", self.config.max_slots);
|
||||
info!(
|
||||
"计算节点已激活,最大并行 Slot 槽位数: {}",
|
||||
self.config.max_slots
|
||||
);
|
||||
|
||||
// Start background heartbeat loop
|
||||
let hb_client = self.client.clone();
|
||||
@@ -71,7 +162,8 @@ impl NodeWorker {
|
||||
if let Ok(mut sys) = s.lock() {
|
||||
sys.refresh_cpu();
|
||||
}
|
||||
}).await;
|
||||
})
|
||||
.await;
|
||||
}
|
||||
sleep(Duration::from_millis(200)).await;
|
||||
{
|
||||
@@ -80,7 +172,8 @@ impl NodeWorker {
|
||||
if let Ok(mut sys) = s.lock() {
|
||||
sys.refresh_cpu();
|
||||
}
|
||||
}).await;
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
loop {
|
||||
@@ -96,13 +189,17 @@ impl NodeWorker {
|
||||
let cpu_usage = sys.global_cpu_info().cpu_usage();
|
||||
let mem_total = sys.total_memory() as f32;
|
||||
let mem_used = sys.used_memory() as f32;
|
||||
let memory_usage = if mem_total > 0.0 { (mem_used / mem_total) * 100.0 } else { 0.0 };
|
||||
let memory_usage = if mem_total > 0.0 {
|
||||
(mem_used / mem_total) * 100.0
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
(cpu_usage, memory_usage)
|
||||
})
|
||||
.await
|
||||
.unwrap_or((0.0, 0.0));
|
||||
|
||||
let active = hb_slots.load(Ordering::Relaxed);
|
||||
let active = hb_slots.load(Ordering::Acquire);
|
||||
let req = NodeHeartbeatRequest {
|
||||
node_id: hb_node_id.clone(),
|
||||
active_slots: active,
|
||||
@@ -110,7 +207,22 @@ impl NodeWorker {
|
||||
memory_usage,
|
||||
};
|
||||
|
||||
let _ = hb_client.post(&hb_url).json(&req).send().await;
|
||||
match hb_client.post(&hb_url).json(&req).send().await {
|
||||
Ok(resp) => {
|
||||
let status = resp.status();
|
||||
// 401/403:token 失效或被吊销。与 claim_task 口径统一:直接退出进程,
|
||||
// 避免心跳线程持续发被拒请求刷日志、占用服务端限流计数。心跳通常比
|
||||
// claim 更高频,往往先于 claim_task 发现吊销。
|
||||
if status.as_u16() == 401 || status.as_u16() == 403 {
|
||||
tracing::error!(
|
||||
"节点 {} 心跳被服务端拒绝 (HTTP {}):node token 已失效或被吊销。请清理 .node_token 文件后重启节点以重新发起注册审批。进程将退出,依赖编排系统重启。",
|
||||
hb_node_id, status
|
||||
);
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
Err(e) => warn!("节点 {} 心跳上报失败: {}", hb_node_id, e),
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
@@ -123,7 +235,7 @@ impl NodeWorker {
|
||||
tokio::spawn(async move {
|
||||
if tokio::signal::ctrl_c().await.is_ok() {
|
||||
info!("收到 Ctrl+C 终止信号,停止领用新任务,准备优雅退出 (再次按 Ctrl+C 可强制立即退出)...");
|
||||
shutdown_signal.store(true, Ordering::SeqCst);
|
||||
shutdown_signal.store(true, Ordering::Release);
|
||||
|
||||
// 二次 Ctrl+C 强行立即退出
|
||||
if tokio::signal::ctrl_c().await.is_ok() {
|
||||
@@ -137,11 +249,11 @@ impl NodeWorker {
|
||||
|
||||
// 带有优雅退出信号响应的任务领用主循环
|
||||
loop {
|
||||
if shutting_down.load(Ordering::Relaxed) {
|
||||
if shutting_down.load(Ordering::Acquire) {
|
||||
break;
|
||||
}
|
||||
|
||||
let active = self.active_slots.load(Ordering::Relaxed);
|
||||
let active = self.active_slots.load(Ordering::Acquire);
|
||||
if (active as usize) < self.config.max_slots {
|
||||
match self.claim_task().await {
|
||||
Ok(Some(task)) => {
|
||||
@@ -149,7 +261,7 @@ impl NodeWorker {
|
||||
info!("与服务端恢复网络连接,已自动重新上线并开始领用计算任务!");
|
||||
was_disconnected = false;
|
||||
}
|
||||
self.active_slots.fetch_add(1, Ordering::SeqCst);
|
||||
self.active_slots.fetch_add(1, Ordering::AcqRel);
|
||||
let client = self.client.clone();
|
||||
let server_url = self.config.server_url.clone();
|
||||
let node_id = self.config.node_id.clone();
|
||||
@@ -158,14 +270,29 @@ impl NodeWorker {
|
||||
let slots_counter = self.active_slots.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let res = execute_task(&client, &server_url, &runtime, &work_dir, &task)
|
||||
.await
|
||||
.map_err(|e| e.to_string());
|
||||
let slot_work_dir = work_dir.join(format!("task_{}", task.task_id));
|
||||
let res =
|
||||
execute_task(&client, &server_url, &runtime, &work_dir, &task)
|
||||
.await
|
||||
.map_err(|e| e.to_string());
|
||||
|
||||
if let Err(e) = report_result(&client, &server_url, &node_id, &task, res).await {
|
||||
let report_res =
|
||||
report_result(&client, &server_url, &node_id, &task, res).await;
|
||||
if report_res.is_ok() {
|
||||
if let Err(e) =
|
||||
crate::executor::cleanup_slot_work_dir(&slot_work_dir).await
|
||||
{
|
||||
warn!(
|
||||
"清理任务 {} 的沙盒目录 {} 失败: {}",
|
||||
task.task_id,
|
||||
slot_work_dir.display(),
|
||||
e
|
||||
);
|
||||
}
|
||||
} else if let Err(ref e) = report_res {
|
||||
warn!("向服务端上报任务 {} 计算结果失败: {}", task.task_id, e);
|
||||
}
|
||||
slots_counter.fetch_sub(1, Ordering::SeqCst);
|
||||
slots_counter.fetch_sub(1, Ordering::AcqRel);
|
||||
});
|
||||
}
|
||||
Ok(None) => {
|
||||
@@ -187,17 +314,17 @@ impl NodeWorker {
|
||||
}
|
||||
|
||||
// 等待在途任务完结(最多等待 30 秒)
|
||||
if self.active_slots.load(Ordering::SeqCst) > 0 {
|
||||
if self.active_slots.load(Ordering::Acquire) > 0 {
|
||||
info!(
|
||||
"正在等待 {} 个在途计算任务优雅完结 (上限 30 秒,按二次 Ctrl+C 可强行中断)...",
|
||||
self.active_slots.load(Ordering::SeqCst)
|
||||
self.active_slots.load(Ordering::Acquire)
|
||||
);
|
||||
}
|
||||
|
||||
let start_wait = std::time::Instant::now();
|
||||
let mut last_log_time = std::time::Instant::now();
|
||||
|
||||
while self.active_slots.load(Ordering::SeqCst) > 0 {
|
||||
while self.active_slots.load(Ordering::Acquire) > 0 {
|
||||
if start_wait.elapsed().as_secs() >= 30 {
|
||||
warn!("在途任务等待超时 (30s),强制退出节点");
|
||||
break;
|
||||
@@ -205,7 +332,7 @@ impl NodeWorker {
|
||||
if last_log_time.elapsed().as_secs() >= 5 {
|
||||
info!(
|
||||
"仍在等待 {} 个在途计算任务完结...",
|
||||
self.active_slots.load(Ordering::SeqCst)
|
||||
self.active_slots.load(Ordering::Acquire)
|
||||
);
|
||||
last_log_time = std::time::Instant::now();
|
||||
}
|
||||
@@ -220,7 +347,21 @@ impl NodeWorker {
|
||||
let claim_url = format!("{}/api/task/claim", self.config.server_url);
|
||||
let resp = self.client.post(&claim_url).send().await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
// 401/403 表明 node token 已被吊销或失效(区别于「暂无任务」与服务端 5xx 故障)。
|
||||
// 服务端故障返回 5xx 会走 !is_success() 的 Ok(None) 分支,仅在网络层/鉴权层拒绝时
|
||||
// 才是真正的吊销。此时继续轮询只会持续产生被拒请求并刷日志,故直接退出进程,
|
||||
// 由编排系统(Docker restart / systemd / k8s)拉起;新进程发现 .node_token 失效后
|
||||
// 会自动走注册审批流程重新申请。
|
||||
if status.as_u16() == 401 || status.as_u16() == 403 {
|
||||
tracing::error!(
|
||||
"领用任务被服务端拒绝 (HTTP {}):node token 已失效或被吊销。请清理 .node_token 文件后重启节点以重新发起注册审批。进程将退出,依赖编排系统重启。",
|
||||
status
|
||||
);
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
if !status.is_success() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
|
||||
@@ -25,5 +25,8 @@ chrono.workspace = true
|
||||
uuid.workspace = true
|
||||
tempfile.workspace = true
|
||||
dotenvy.workspace = true
|
||||
sha2.workspace = true
|
||||
hex.workspace = true
|
||||
subtle = "2"
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
//! 管理 API(Admin 角色)。
|
||||
//!
|
||||
//! 提供 node 凭据的可视化与运维操作,供 Dashboard 管理界面调用:
|
||||
//! - 列出所有节点及其凭据状态(在线/token 是否有效/吊销/颁发时间)
|
||||
//! - 吊销指定节点的专属 token(立即失效,不影响其他节点)
|
||||
//! - 重新颁发指定节点的专属 token(返回新明文,旧 token 失效)
|
||||
//!
|
||||
//! 这些端点均要求 Admin 角色(见 mod.rs 授权矩阵),node 自身无权操作他人或自身凭据,
|
||||
//! 从而保证「吊销/重发」是管理员主动行为,避免被攻陷节点篡改凭据体系。
|
||||
|
||||
use super::{is_valid_node_id, AppState};
|
||||
use axum::{
|
||||
extract::{Path as AxumPath, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// GET /api/admin/nodes — 列出全部节点及凭据状态。
|
||||
pub async fn list_nodes(
|
||||
State(state): State<AppState>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
match state.db.list_nodes_with_credentials().await {
|
||||
Ok(list) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({ "success": true, "message": "成功获取节点列表", "data": list })),
|
||||
)),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /api/admin/nodes/:node_id/revoke — 吊销指定节点的专属 token。
|
||||
///
|
||||
/// 吊销后该 node 的现有 token 立即失效,须重新走注册流程领取新 token。
|
||||
/// 操作幂等:对无凭据记录或已吊销的节点调用不会报错。
|
||||
pub async fn revoke_node(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(node_id): AxumPath<String>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// node_id 白名单校验,防止注入或异常输入(与 register_node 的 node_id 来源口径一致)
|
||||
if !is_valid_node_id(&node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID 参数".to_string(),
|
||||
));
|
||||
}
|
||||
match state.db.revoke_node_token(&node_id).await {
|
||||
Ok(_) => {
|
||||
info!("管理员已吊销节点 {} 的专属 token", node_id);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(
|
||||
json!({ "success": true, "message": format!("节点 '{}' 的 token 已吊销", node_id) }),
|
||||
),
|
||||
))
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("吊销节点 {} token 失败: {}", node_id, e);
|
||||
Err(e.into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /api/admin/nodes/:node_id/reissue — 重新颁发指定节点的专属 token。
|
||||
///
|
||||
/// 旧 token 立即失效,返回新 token 明文(仅此一次,DB 只存 hash)。
|
||||
/// 节点需用新 token 重新注册或由管理员手动同步到节点本地 `.node_token`。
|
||||
pub async fn reissue_node(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(node_id): AxumPath<String>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
if !is_valid_node_id(&node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID 参数".to_string(),
|
||||
));
|
||||
}
|
||||
// 仅允许对已注册的节点重发 token(防止凭据表被写入幽灵 node_id)
|
||||
match state.db.get_node_exists(&node_id).await {
|
||||
Ok(false) => {
|
||||
return Err(crate::api::AppError::NotFound(format!(
|
||||
"节点 '{}' 不存在,请先注册",
|
||||
node_id
|
||||
)));
|
||||
}
|
||||
Ok(true) => {}
|
||||
Err(e) => return Err(e.into()),
|
||||
}
|
||||
match state.db.issue_node_token(&node_id).await {
|
||||
Ok(new_token) => {
|
||||
info!("管理员已为节点 {} 重新颁发专属 token", node_id);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"success": true,
|
||||
"message": format!("节点 '{}' 的 token 已重新颁发,请将新 token 同步到该节点", node_id),
|
||||
"node_token": new_token,
|
||||
})),
|
||||
))
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("为节点 {} 重新颁发 token 失败: {}", node_id, e);
|
||||
Err(e.into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /api/admin/nodes/:node_id/approve — 管理员同意节点接入申请。
|
||||
pub async fn approve_node(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(node_id): AxumPath<String>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
if !is_valid_node_id(&node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID 参数".to_string(),
|
||||
));
|
||||
}
|
||||
match state.db.approve_node(&node_id).await {
|
||||
Ok(_token) => {
|
||||
info!("管理员已同意节点 {} 的接入申请并生成专属 Token", node_id);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(
|
||||
json!({ "success": true, "message": format!("节点 '{}' 已授权加入集群", node_id) }),
|
||||
),
|
||||
))
|
||||
}
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /api/admin/nodes/:node_id/reject — 管理员拒绝节点接入申请。
|
||||
pub async fn reject_node(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(node_id): AxumPath<String>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
if !is_valid_node_id(&node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID 参数".to_string(),
|
||||
));
|
||||
}
|
||||
match state.db.reject_node(&node_id).await {
|
||||
Ok(_) => {
|
||||
info!("管理员已拒绝节点 {} 的接入申请并移除", node_id);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(
|
||||
json!({ "success": true, "message": format!("已拒绝节点 '{}' 的接入申请", node_id) }),
|
||||
),
|
||||
))
|
||||
}
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
//! 管理员表单登录与凭据校验 API。
|
||||
//!
|
||||
//! 提供基于短密码的身份认证服务:
|
||||
//! - POST /api/login:校验管理员密码,成功后返回 Admin Token,并记录 IP 错误次数防止暴力破解。
|
||||
//! - GET /api/auth/check:由 auth_middleware 保护,供前端初始化时检测当前保存的 Token 是否有效。
|
||||
|
||||
use super::{ct_eq_str, AppState};
|
||||
use axum::{
|
||||
extract::{ConnectInfo, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::net::SocketAddr;
|
||||
use tracing::{info, warn};
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct LoginRequest {
|
||||
pub password: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct LoginResponse {
|
||||
pub success: bool,
|
||||
pub message: String,
|
||||
pub token: Option<String>,
|
||||
}
|
||||
|
||||
/// POST /api/login — 管理员密码登录端点。
|
||||
pub async fn login(
|
||||
State(state): State<AppState>,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
Json(req): Json<LoginRequest>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
let client_ip = addr.ip();
|
||||
|
||||
// 限流检查:5 分钟内最多允许 5 次失败尝试(基于 RateLimiter 防暴力破解)
|
||||
if state.rate_limiter.is_rate_limited(client_ip) {
|
||||
warn!("客户端 IP {} 登录失败次数过多,已临时封禁锁定", client_ip);
|
||||
return Err(crate::api::AppError::TooManyRequests(
|
||||
"登录失败次数过多,已被临时锁定,请 5 分钟后再试".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let admin_token = match state.admin_token.as_deref() {
|
||||
Some(t) if !t.is_empty() => t,
|
||||
_ => {
|
||||
warn!("系统未配置 admin_token 且鉴权未禁用,拒绝登录");
|
||||
return Err(crate::api::AppError::Forbidden(
|
||||
"服务端未配置管理员凭据,请检查配置文件".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// 恒定时间密码比对(防时序旁路攻击)
|
||||
if ct_eq_str(&req.password, admin_token) {
|
||||
// 生成随机 64 位 Session Token
|
||||
let session_token = format!(
|
||||
"{}{}",
|
||||
uuid::Uuid::new_v4().simple(),
|
||||
uuid::Uuid::new_v4().simple()
|
||||
);
|
||||
let expiry = std::time::Instant::now() + std::time::Duration::from_secs(24 * 3600);
|
||||
|
||||
// 存储 Token 到内存中(带容量上限清理)
|
||||
{
|
||||
let mut sessions = state.admin_sessions.write().await;
|
||||
let now = std::time::Instant::now();
|
||||
// 1. 清理已过期的 session
|
||||
sessions.retain(|_, exp| *exp > now);
|
||||
// 2. 若超出容量限制,淘汰最老/最快过期的 session
|
||||
while sessions.len() >= crate::api::MAX_ADMIN_SESSIONS {
|
||||
if let Some(oldest_key) = sessions
|
||||
.iter()
|
||||
.min_by_key(|(_, exp)| **exp)
|
||||
.map(|(k, _)| k.clone())
|
||||
{
|
||||
sessions.remove(&oldest_key);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
sessions.insert(session_token.clone(), expiry);
|
||||
}
|
||||
|
||||
info!(
|
||||
"客户端 IP {} 密码验证成功,已颁发 Admin Session Token",
|
||||
client_ip
|
||||
);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(LoginResponse {
|
||||
success: true,
|
||||
message: "登录成功".to_string(),
|
||||
token: Some(session_token),
|
||||
}),
|
||||
))
|
||||
} else {
|
||||
warn!("客户端 IP {} 登录密码校验失败", client_ip);
|
||||
// 记录一次失败
|
||||
state.rate_limiter.record_failure(client_ip);
|
||||
Err(crate::api::AppError::Unauthorized(
|
||||
"管理员密码错误,请重新输入".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// GET /api/auth/check — 校验当前 Admin Token 是否有效。
|
||||
///
|
||||
/// 放在 auth_middleware(Role::Admin)之后,只要到达此 handler 说明 Token 校验必定成功。
|
||||
pub async fn check_auth() -> impl IntoResponse {
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(serde_json::json!({
|
||||
"success": true,
|
||||
"message": "Token 验证有效",
|
||||
"authenticated": true
|
||||
})),
|
||||
)
|
||||
}
|
||||
@@ -1,37 +1,81 @@
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::Path as AxumPath,
|
||||
http::{header, StatusCode},
|
||||
response::IntoResponse,
|
||||
};
|
||||
use axum::{body::Body, extract::Path as AxumPath, http::header, response::IntoResponse};
|
||||
use std::path::{Path, PathBuf};
|
||||
use tokio::fs::File;
|
||||
use tokio_util::io::ReaderStream;
|
||||
|
||||
pub async fn download_single_data_file(AxumPath(filename): AxumPath<String>) -> axum::response::Response {
|
||||
pub async fn download_single_data_file(
|
||||
AxumPath(filename): AxumPath<String>,
|
||||
) -> Result<axum::response::Response, crate::api::AppError> {
|
||||
let safe_name = Path::new(&filename)
|
||||
.file_name()
|
||||
.map(|s| s.to_string_lossy().to_string())
|
||||
.unwrap_or_default();
|
||||
|
||||
if safe_name.is_empty() || safe_name.starts_with('.') {
|
||||
return (StatusCode::BAD_REQUEST, "无效的数据文件名").into_response();
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"无效的数据文件名".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// 严苛白名单过滤:严防 `..`、特殊符号注入及路径穿透攻击
|
||||
if !safe_name.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@') {
|
||||
tracing::warn!("拦截到疑似非法字符构造的敏感及越界资源抓取行为: {}", safe_name);
|
||||
return (StatusCode::BAD_REQUEST, "参数非法,请求的文件包含系统不许可的危险专属占位或路径重定向字符").into_response();
|
||||
if !safe_name.chars().all(|c| {
|
||||
c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@'
|
||||
}) {
|
||||
tracing::warn!(
|
||||
"拦截到疑似非法字符构造的敏感及越界资源抓取行为: {}",
|
||||
safe_name
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"参数非法,请求的文件包含系统不许可的危险专属占位或路径重定向字符".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let rel_path = format!("assets/data/{}", safe_name);
|
||||
tracing::debug!("服务端处理数据文件下载请求: {}", safe_name);
|
||||
stream_asset_file(&rel_path, "application/octet-stream").await.into_response()
|
||||
stream_asset_file(&rel_path, "application/octet-stream").await
|
||||
}
|
||||
|
||||
pub async fn download_linelist() -> axum::response::Response {
|
||||
let linelist_path = std::env::var("DCTS_LINELIST_PATH").unwrap_or_else(|_| "assets/gfVIS99.dat".to_string());
|
||||
stream_asset_file(&linelist_path, "application/octet-stream").await.into_response()
|
||||
pub async fn download_linelist() -> Result<axum::response::Response, crate::api::AppError> {
|
||||
let linelist_path =
|
||||
std::env::var("DCTS_LINELIST_PATH").unwrap_or_else(|_| "assets/gfVIS99.dat".to_string());
|
||||
|
||||
// 路径规约校验:DCTS_LINELIST_PATH 解析后的绝对路径必须落在 assets 根目录内,
|
||||
// 防止环境变量被设为 ../../etc/passwd 之类导致任意文件流出。
|
||||
// assets 根目录优先取 DCTS_ASSETS_DIR,回退到相对路径 assets。
|
||||
let assets_root = std::env::var("DCTS_ASSETS_DIR").unwrap_or_else(|_| "assets".to_string());
|
||||
if !is_path_within_assets(&linelist_path, &assets_root) {
|
||||
tracing::warn!(
|
||||
"DCTS_LINELIST_PATH '{}' 不在 assets 根目录 '{}' 内,拒绝下载",
|
||||
linelist_path,
|
||||
assets_root
|
||||
);
|
||||
return Err(crate::api::AppError::Forbidden(
|
||||
"请求的谱线文件路径越界,已被拒绝".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
stream_asset_file(&linelist_path, "application/octet-stream").await
|
||||
}
|
||||
|
||||
/// 校验 target 路径(经 canonicalize 后)是否落在 assets 根目录之内。
|
||||
/// 对不存在的路径(canonicalize 失败)回退到 starts_with 的词法比较,宁可偏严。
|
||||
fn is_path_within_assets(target: &str, assets_root: &str) -> bool {
|
||||
// 严防 `..` 词法穿透
|
||||
if target.contains("..") {
|
||||
return false;
|
||||
}
|
||||
|
||||
let target_path = Path::new(target);
|
||||
let root_path = Path::new(assets_root);
|
||||
|
||||
let target_abs = std::fs::canonicalize(target_path).ok();
|
||||
let root_abs = std::fs::canonicalize(root_path).ok();
|
||||
|
||||
match (target_abs, root_abs) {
|
||||
(Some(t), Some(r)) => t.starts_with(&r),
|
||||
// 路径尚未存在时用词法前缀比较(canonicalize 需要文件存在)
|
||||
_ => target_path.starts_with(root_path),
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_asset(rel_path: &str) -> Option<PathBuf> {
|
||||
@@ -65,10 +109,17 @@ fn resolve_asset(rel_path: &str) -> Option<PathBuf> {
|
||||
None
|
||||
}
|
||||
|
||||
async fn stream_asset_file(rel_path: &str, content_type: &'static str) -> impl IntoResponse {
|
||||
async fn stream_asset_file(
|
||||
rel_path: &str,
|
||||
content_type: &'static str,
|
||||
) -> Result<axum::response::Response, crate::api::AppError> {
|
||||
let resolved_path = match resolve_asset(rel_path) {
|
||||
Some(p) => p,
|
||||
None => return (StatusCode::NOT_FOUND, "资源数据文件不存在").into_response(),
|
||||
None => {
|
||||
return Err(crate::api::AppError::NotFound(
|
||||
"资源数据文件不存在".to_string(),
|
||||
))
|
||||
}
|
||||
};
|
||||
|
||||
match File::open(&resolved_path).await {
|
||||
@@ -89,9 +140,11 @@ async fn stream_asset_file(rel_path: &str, content_type: &'static str) -> impl I
|
||||
(header::CONTENT_DISPOSITION, disposition),
|
||||
];
|
||||
|
||||
(headers, body).into_response()
|
||||
Ok((headers, body).into_response())
|
||||
}
|
||||
Err(e) => {
|
||||
let boxed_err: anyhow::Error = e.into();
|
||||
Err(crate::api::AppError::Internal(boxed_err))
|
||||
}
|
||||
Err(_) => (StatusCode::INTERNAL_SERVER_ERROR, "无法读取资源数据文件").into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
use axum::{
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tracing::error;
|
||||
|
||||
pub enum AppError {
|
||||
BadRequest(String),
|
||||
Unauthorized(String),
|
||||
Forbidden(String),
|
||||
NotFound(String),
|
||||
Conflict(String),
|
||||
TooManyRequests(String),
|
||||
Internal(anyhow::Error),
|
||||
}
|
||||
|
||||
impl IntoResponse for AppError {
|
||||
fn into_response(self) -> Response {
|
||||
let (status, error_message) = match self {
|
||||
AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg),
|
||||
AppError::Unauthorized(msg) => (StatusCode::UNAUTHORIZED, msg),
|
||||
AppError::Forbidden(msg) => (StatusCode::FORBIDDEN, msg),
|
||||
AppError::NotFound(msg) => (StatusCode::NOT_FOUND, msg),
|
||||
AppError::Conflict(msg) => (StatusCode::CONFLICT, msg),
|
||||
AppError::TooManyRequests(msg) => (StatusCode::TOO_MANY_REQUESTS, msg),
|
||||
AppError::Internal(err) => {
|
||||
error!("Internal server error: {:?}", err);
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"Internal server error".to_string(),
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
let body = Json(json!({
|
||||
"success": false,
|
||||
"message": error_message,
|
||||
"data": serde_json::Value::Null
|
||||
}));
|
||||
|
||||
(status, body).into_response()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<anyhow::Error> for AppError {
|
||||
fn from(inner: anyhow::Error) -> Self {
|
||||
AppError::Internal(inner)
|
||||
}
|
||||
}
|
||||
+249
-41
@@ -1,20 +1,28 @@
|
||||
pub mod admin;
|
||||
pub mod auth;
|
||||
pub mod data;
|
||||
pub mod error;
|
||||
pub mod node;
|
||||
pub mod rate_limit;
|
||||
pub mod seed;
|
||||
pub mod status;
|
||||
pub mod task;
|
||||
pub mod workflow;
|
||||
|
||||
pub use error::AppError;
|
||||
|
||||
use crate::db::Database;
|
||||
use crate::scheduler::GridScheduler;
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::{header, Request, StatusCode},
|
||||
middleware::Next,
|
||||
response::IntoResponse,
|
||||
};
|
||||
use crate::db::Database;
|
||||
use crate::scheduler::GridScheduler;
|
||||
use mq::sqlite_queue::SqliteTaskQueue;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::sync::Arc;
|
||||
use subtle::ConstantTimeEq;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
@@ -22,53 +30,253 @@ pub struct AppState {
|
||||
pub queue: Arc<SqliteTaskQueue>,
|
||||
pub scheduler: Arc<GridScheduler>,
|
||||
pub results_dir: String,
|
||||
/// 限流与密码防暴破限速器
|
||||
pub rate_limiter: rate_limit::RateLimiter,
|
||||
/// 兼容字段:Some 表示「已启用某种鉴权」,用于 main.rs 决定是否挂载鉴权中间件。
|
||||
pub auth_token: Option<String>,
|
||||
/// Admin 凭据(管理 Dashboard / workflow 写操作)。
|
||||
pub admin_token: Option<String>,
|
||||
/// 应急开关:跳过全部鉴权(仅本地调试)。
|
||||
pub auth_disabled: bool,
|
||||
/// 动态 Session Tokens(登录后发放),设置 24 小时过期
|
||||
pub admin_sessions:
|
||||
Arc<tokio::sync::RwLock<std::collections::HashMap<String, std::time::Instant>>>,
|
||||
}
|
||||
|
||||
/// 固定时间敏感字符串一致性核验函数,彻底消解时序测信道猜测危险
|
||||
fn constant_time_eq(a: &str, b: &str) -> bool {
|
||||
let a_bytes = a.as_bytes();
|
||||
let b_bytes = b.as_bytes();
|
||||
let mut diff = (a_bytes.len() ^ b_bytes.len()) as u64;
|
||||
// 遍历目标 secret (b_bytes) 的完整长度,使耗时仅受 server 预期 token 长度决定
|
||||
for (i, &y) in b_bytes.iter().enumerate() {
|
||||
let x = if i < a_bytes.len() { a_bytes[i] } else { 0 };
|
||||
diff |= (x ^ y) as u64;
|
||||
/// Admin Session 最大保存上限
|
||||
pub const MAX_ADMIN_SESSIONS: usize = 100;
|
||||
|
||||
/// 已认证的 Node 身份(中间件校验 node token 通过后注入 request extension)。
|
||||
///
|
||||
/// 下游 handler(heartbeat / claim / report)通过 `Extension<AuthenticatedNode>` 取出,
|
||||
/// 用于校验请求体里声称的 node_id 与 token 绑定的 node_id 一致,杜绝跨节点冒充。
|
||||
#[derive(Clone)]
|
||||
pub struct AuthenticatedNode {
|
||||
pub node_id: String,
|
||||
}
|
||||
|
||||
/// 授权角色:决定某条路径需要哪类主体才能访问。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum Role {
|
||||
/// 公开免鉴权端点:登录 /login,节点注册申请 /node/register,审批状态查询 /node/check_status。
|
||||
Public,
|
||||
/// Node 运行态:心跳/领任务/上报/下载种子与数据。需 node 专属 token。
|
||||
Node,
|
||||
/// 管理操作:workflow CRUD、起停、查看 status、审批节点。需 admin token。
|
||||
Admin,
|
||||
}
|
||||
|
||||
/// 路径 → 角色授权矩阵。
|
||||
///
|
||||
/// 设计依据(最小权限):
|
||||
/// - Admin 写操作(workflow CRUD / start / stop / status / approve / reject)只对 admin token 开放。
|
||||
/// - Node 运行态接口只认 node 专属 token(管理员在 Dashboard 审批后颁发,绑定 node_id,可吊销)。
|
||||
/// - 注册端点 /node/register 和状态轮询 /node/check_status 为 Public 免凭据(提交申请 ➔ 待管理员审批)。
|
||||
///
|
||||
/// 注意:路径已去掉 `/api` 前缀(nest 挂载后中间件看到的 path 不含 nest 前缀)。
|
||||
fn required_role(path: &str, method: &axum::http::Method) -> Option<Role> {
|
||||
use axum::http::Method;
|
||||
// 公开免鉴权端点
|
||||
if (path == "/login" || path == "/node/register" || path == "/node/check_status")
|
||||
&& method == Method::POST
|
||||
{
|
||||
return Some(Role::Public);
|
||||
}
|
||||
diff == 0
|
||||
// 校验身份与状态 -> Admin
|
||||
if path == "/auth/check" && method == Method::GET {
|
||||
return Some(Role::Admin);
|
||||
}
|
||||
// 写操作 → Admin
|
||||
if path == "/workflows" && (method == Method::POST || method == Method::GET) {
|
||||
return Some(Role::Admin);
|
||||
}
|
||||
if path.starts_with("/workflows/") {
|
||||
// GET/PUT/DELETE /workflows/:name, POST /start|stop → Admin
|
||||
return Some(Role::Admin);
|
||||
}
|
||||
if path == "/status" && method == Method::GET {
|
||||
return Some(Role::Admin);
|
||||
}
|
||||
// 管理 API(节点凭据查看/审批/吊销/重发)→ Admin
|
||||
if path.starts_with("/admin/") {
|
||||
return Some(Role::Admin);
|
||||
}
|
||||
// Node 运行态 → Node
|
||||
if path == "/node/heartbeat" && method == Method::POST {
|
||||
return Some(Role::Node);
|
||||
}
|
||||
if path == "/task/claim" && method == Method::POST {
|
||||
return Some(Role::Node);
|
||||
}
|
||||
if path == "/task/report" && method == Method::POST {
|
||||
return Some(Role::Node);
|
||||
}
|
||||
if path.starts_with("/seed/") && method == Method::GET {
|
||||
return Some(Role::Node);
|
||||
}
|
||||
if (path.starts_with("/data/file/") || path == "/data/linelist") && method == Method::GET {
|
||||
return Some(Role::Node);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Axum 鉴权中间件:若 AppState 中配置了 auth_token 则强制校验 Bearer Token 或 X-API-Key
|
||||
pub async fn auth_middleware(
|
||||
State(state): State<AppState>,
|
||||
req: Request<axum::body::Body>,
|
||||
next: Next,
|
||||
) -> impl IntoResponse {
|
||||
if let Some(ref expected_token) = state.auth_token {
|
||||
let auth_header = req
|
||||
.headers()
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok());
|
||||
let api_key_header = req
|
||||
.headers()
|
||||
.get("x-api-key")
|
||||
.and_then(|v| v.to_str().ok());
|
||||
/// 恒定时间字符串比较。
|
||||
///
|
||||
/// 先对两个串各自做 SHA-256,再比较等长摘要(32 字节),彻底消除长度时序旁路——
|
||||
/// 任意长度的输入都产生相同长度的摘要,比较耗时固定,攻击者无法通过响应时间探得 token 长度。
|
||||
fn ct_eq_str(a: &str, b: &str) -> bool {
|
||||
let ha = {
|
||||
let mut h = Sha256::new();
|
||||
h.update(a.as_bytes());
|
||||
h.finalize()
|
||||
};
|
||||
let hb = {
|
||||
let mut h = Sha256::new();
|
||||
h.update(b.as_bytes());
|
||||
h.finalize()
|
||||
};
|
||||
ha.ct_eq(&hb).into()
|
||||
}
|
||||
|
||||
let token_valid = match (auth_header, api_key_header) {
|
||||
(Some(auth), _) if auth.starts_with("Bearer ") => constant_time_eq(&auth[7..], expected_token),
|
||||
(Some(auth), _) => constant_time_eq(auth, expected_token),
|
||||
(_, Some(key)) => constant_time_eq(key, expected_token),
|
||||
_ => false,
|
||||
};
|
||||
/// node_id 白名单:字母、数字、点、下划线、连字符,长度 1-128。
|
||||
/// 用于 register_node / admin revoke / reissue 统一入口校验,与 Dashboard XSS 防护口径一致。
|
||||
pub(crate) fn is_valid_node_id(id: &str) -> bool {
|
||||
!id.is_empty()
|
||||
&& id.len() <= 128
|
||||
&& id
|
||||
.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-')
|
||||
}
|
||||
|
||||
if !token_valid {
|
||||
return (
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"Unauthorized: Invalid or missing authentication token",
|
||||
)
|
||||
.into_response();
|
||||
/// host_name 白名单:可打印 ASCII(排除控制字符),长度 1-128。
|
||||
/// 防止 host_name 携带 HTML/控制字符进入管理 Dashboard 触发存储型 XSS 或污染显示。
|
||||
pub(crate) fn is_valid_host_name(name: &str) -> bool {
|
||||
!name.is_empty()
|
||||
&& name.len() <= 128
|
||||
&& name.chars().all(|c| c.is_ascii() && !c.is_ascii_control())
|
||||
}
|
||||
|
||||
/// 从请求头提取凭据原文(支持 `Authorization: Bearer <t>` 与 `X-API-Key: <t>`)。
|
||||
///
|
||||
/// 安全:非 `Bearer ` 前缀的 Authorization 一律视为无 token(不再回退为裸头值比较),
|
||||
/// 避免 `Authorization: Basic ...` 之类的上游代理头被误送入 token 比对。
|
||||
fn extract_token(req: &Request<axum::body::Body>) -> Option<String> {
|
||||
if let Some(auth) = req
|
||||
.headers()
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
if let Some(rest) = auth.strip_prefix("Bearer ") {
|
||||
if !rest.is_empty() {
|
||||
return Some(rest.to_string());
|
||||
}
|
||||
}
|
||||
// 非 Bearer 前缀或空值:不作为 token
|
||||
}
|
||||
if let Some(key) = req.headers().get("x-api-key").and_then(|v| v.to_str().ok()) {
|
||||
if !key.is_empty() {
|
||||
return Some(key.to_string());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Axum 鉴权中间件(L2)。
|
||||
///
|
||||
/// 流程:
|
||||
/// 1. 应急关闭(auth_disabled)→ 直接放行。
|
||||
/// 2. 路径不在授权矩阵 → 视为未公开接口,拒绝(401)。
|
||||
/// 3. 按角色校验对应凭据:
|
||||
/// - Admin: admin token 恒定时间比对。
|
||||
/// - Node: node 专属 token 经 DB 反查 node_id(token 只存 hash)。
|
||||
/// 4. Node 角色额外校验:请求声称的 node_id 须与 token 绑定的 node_id 一致
|
||||
/// (防 A 节点用 B 节点的 token 越权操作)。claim_task / data 下载无 node_id
|
||||
/// 输入,则仅校验 token 有效即可。
|
||||
pub async fn auth_middleware(
|
||||
State(state): State<AppState>,
|
||||
mut req: Request<axum::body::Body>,
|
||||
next: Next,
|
||||
) -> impl IntoResponse {
|
||||
// 应急关闭:本地调试专用,绕过全部校验
|
||||
if state.auth_disabled {
|
||||
return next.run(req).await.into_response();
|
||||
}
|
||||
|
||||
let path = req.uri().path().to_string();
|
||||
let method = req.method().clone();
|
||||
let role = match required_role(&path, &method) {
|
||||
Some(Role::Public) => {
|
||||
return next.run(req).await.into_response();
|
||||
}
|
||||
Some(r) => r,
|
||||
None => {
|
||||
// 未在矩阵中的路径一律拒绝(默认拒绝原则)
|
||||
return (StatusCode::UNAUTHORIZED, "Unauthorized").into_response();
|
||||
}
|
||||
};
|
||||
|
||||
let token = match extract_token(&req) {
|
||||
Some(t) => t,
|
||||
None => {
|
||||
return (StatusCode::UNAUTHORIZED, "Unauthorized: missing token").into_response();
|
||||
}
|
||||
};
|
||||
|
||||
// 审计日志:仅记录写操作(POST/PUT/DELETE)的「谁、做了什么」,不记请求体(防泄露)。
|
||||
// 在校验通过后记录 subject;校验失败由 401 分支体现,不单独审计。
|
||||
use axum::http::Method;
|
||||
let is_write = matches!(method, Method::POST | Method::PUT | Method::DELETE);
|
||||
|
||||
match role {
|
||||
Role::Public => unreachable!(),
|
||||
Role::Admin => {
|
||||
let mut valid = false;
|
||||
if let Some(ref admin) = state.admin_token {
|
||||
if ct_eq_str(&token, admin) {
|
||||
valid = true;
|
||||
}
|
||||
}
|
||||
if !valid {
|
||||
let mut sessions = state.admin_sessions.write().await;
|
||||
let now = std::time::Instant::now();
|
||||
sessions.retain(|_, expiry| *expiry > now);
|
||||
if sessions.contains_key(&token) {
|
||||
valid = true;
|
||||
}
|
||||
}
|
||||
|
||||
if valid {
|
||||
if is_write {
|
||||
tracing::info!(target: "dcts_audit", "AUDIT subject=admin method={} path={}", method, path);
|
||||
}
|
||||
return next.run(req).await.into_response();
|
||||
}
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"Unauthorized: invalid admin token",
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
Role::Node => {
|
||||
// 用 token 反查所属 node_id(DB 只存 hash,明文不落库)
|
||||
match state.db.find_node_by_token(&token).await {
|
||||
Some(token_node_id) => {
|
||||
// 仅对非例行高频请求(如任务结果上报 /task/report)记录 AUDIT 审计日志,
|
||||
// 成功的例行心跳 (/node/heartbeat) 与空闲领任务 (/task/claim) 静默跳过。
|
||||
if is_write && path != "/node/heartbeat" && path != "/task/claim" {
|
||||
tracing::info!(target: "dcts_audit", "AUDIT subject=node:{} method={} path={}", token_node_id, method, path);
|
||||
}
|
||||
req.extensions_mut().insert(AuthenticatedNode {
|
||||
node_id: token_node_id,
|
||||
});
|
||||
next.run(req).await.into_response()
|
||||
}
|
||||
None => (
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"Unauthorized: invalid or revoked node token",
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
next.run(req).await.into_response()
|
||||
}
|
||||
|
||||
@@ -1,25 +1,152 @@
|
||||
use super::AppState;
|
||||
use axum::{extract::State, response::IntoResponse, Json};
|
||||
use super::{is_valid_host_name, is_valid_node_id, AppState, AuthenticatedNode};
|
||||
use axum::{
|
||||
extract::{Extension, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
use common::models::{NodeHeartbeatRequest, NodeRegisterRequest};
|
||||
use serde_json::json;
|
||||
use tracing::{info, warn};
|
||||
|
||||
pub async fn register_node(
|
||||
State(state): State<AppState>,
|
||||
auth_node: Option<Extension<AuthenticatedNode>>,
|
||||
Json(req): Json<NodeRegisterRequest>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// 入口白名单校验
|
||||
if !is_valid_node_id(&req.node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID(仅允许字母、数字、点、下划线、连字符,长度 1-128)".to_string(),
|
||||
));
|
||||
}
|
||||
if !is_valid_host_name(&req.host_name) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的主机名(仅允许可打印 ASCII,长度 1-128)".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// 已认证已拿到 Token 的节点刷新元数据配置
|
||||
if let Some(Extension(ref auth)) = auth_node {
|
||||
if auth.node_id == req.node_id {
|
||||
let _ = state.db.register_node(&req).await;
|
||||
info!("已授权节点 {} 刷新配置成功", req.node_id);
|
||||
return Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "approved",
|
||||
"message": "节点配置更新成功",
|
||||
"node_token": null,
|
||||
})),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// 申请注册新节点(免凭据提交申请,进入 pending_approval 状态)
|
||||
match state.db.register_node(&req).await {
|
||||
Ok(_) => Json(json!({"status": "ok", "message": "节点注册成功"})),
|
||||
Err(e) => Json(json!({"status": "error", "message": e.to_string()})),
|
||||
Ok(true) => {
|
||||
info!(
|
||||
"接收到新节点 {} 的注册申请,已加入待审批 (pending_approval) 队列",
|
||||
req.node_id
|
||||
);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "pending_approval",
|
||||
"message": "节点注册申请已成功提交!请在管理 Dashboard 控制台上点击【同意接入】授权该节点",
|
||||
"node_token": null,
|
||||
})),
|
||||
))
|
||||
}
|
||||
Ok(false) => {
|
||||
// 节点已处于待审批或已存在列表
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "pending_approval",
|
||||
"message": "节点注册申请等待管理员审批中",
|
||||
"node_token": null,
|
||||
})),
|
||||
))
|
||||
}
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
pub struct CheckNodeStatusRequest {
|
||||
pub node_id: String,
|
||||
}
|
||||
|
||||
/// POST /api/node/check_status — Node 端轮询检查审批结果。
|
||||
pub async fn check_node_status(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<CheckNodeStatusRequest>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
if !is_valid_node_id(&req.node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID 参数".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// 尝试拉取取走即焚的暂存明文 Token
|
||||
match state.db.take_pending_node_token(&req.node_id).await {
|
||||
Ok(Some(raw_token)) => {
|
||||
info!(
|
||||
"节点 {} 的注册申请已被管理员审批同意,下发专属 Token",
|
||||
req.node_id
|
||||
);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "approved",
|
||||
"message": "节点已通过审批授权",
|
||||
"node_token": raw_token,
|
||||
})),
|
||||
))
|
||||
}
|
||||
Ok(None) | Err(_) => {
|
||||
// 查节点表状态
|
||||
match state.db.get_node_exists(&req.node_id).await {
|
||||
Ok(true) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "pending_approval",
|
||||
"message": "等待管理员在控制台点击同意",
|
||||
"node_token": null,
|
||||
})),
|
||||
)),
|
||||
_ => Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "rejected",
|
||||
"message": "节点注册申请未通过或已被移除",
|
||||
"node_token": null,
|
||||
})),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn heartbeat_node(
|
||||
State(state): State<AppState>,
|
||||
Extension(auth_node): Extension<AuthenticatedNode>,
|
||||
Json(req): Json<NodeHeartbeatRequest>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// 身份绑定校验:请求体声称的 node_id 必须与 token 绑定的 node_id 一致,
|
||||
// 杜绝「持有 A 节点 token 却冒充 B 节点发心跳」的跨节点越权。
|
||||
if req.node_id != auth_node.node_id {
|
||||
warn!(
|
||||
"节点心跳身份校验失败:token 绑定 node={},但请求体声称 node_id={}",
|
||||
auth_node.node_id, req.node_id
|
||||
);
|
||||
return Err(crate::api::AppError::Forbidden(
|
||||
"node_id 与凭据不匹配".to_string(),
|
||||
));
|
||||
}
|
||||
match state.db.heartbeat_node(&req).await {
|
||||
Ok(_) => Json(json!({"status": "ok"})),
|
||||
Err(e) => Json(json!({"status": "error", "message": e.to_string()})),
|
||||
Ok(_) => Ok(Json(json!({"status": "ok"}))),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,196 @@
|
||||
//! 鉴权失败速率限制中间件(防 token 在线暴力)。
|
||||
//!
|
||||
//! 设计:对返回 401 的请求按客户端 IP 维护滑动窗口失败计数。当某 IP 在窗口内
|
||||
//! 累计失败超过阈值,后续请求直接返回 429(持续到窗口内计数回落)。
|
||||
//!
|
||||
//! 仅作用于鉴权路径(与 auth_middleware 叠加),不影响已认证的正常业务流。
|
||||
//! 已认证请求返回 2xx,不计入失败窗口,因此合法节点/管理员的高频调用不受影响。
|
||||
//!
|
||||
//! IP 来源:优先取 `X-Forwarded-For` 首段(反代场景),回退到连接的 `ConnectInfo<SocketAddr>`
|
||||
//! (需 main.rs 用 `into_make_service_with_connect_info` 启动)。两者都拿不到时按"未知 IP"聚合。
|
||||
|
||||
use axum::{
|
||||
extract::{ConnectInfo, State},
|
||||
http::Request,
|
||||
middleware::Next,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
use tracing::warn;
|
||||
|
||||
/// 限流状态:按 IP 维护近窗口内的失败时间戳队列。
|
||||
#[derive(Clone)]
|
||||
pub struct RateLimiter {
|
||||
inner: Arc<Mutex<HashMap<IpAddr, VecDeque<Instant>>>>,
|
||||
window: Duration,
|
||||
max_failures: usize,
|
||||
/// 计数策略:
|
||||
/// - `false`(默认,通用 API 限流器):仅对鉴权失败(400/401/403)的响应计数。
|
||||
/// - `true`(注册端点专用限流器):对匹配路径(如 `/node/register`)的**所有**响应计数,
|
||||
/// 无论成败——这是对注册接口的独立节流设计,防止恶意频繁注册。
|
||||
///
|
||||
/// 历史问题:此前中间件对所有 `/node/register` 请求无条件计数,导致该 limiter 若复用为
|
||||
/// 通用 API 限流器时,20 次成功注册会把整个 IP 锁出所有 `/api/*` 端点(跨端点连锁)。
|
||||
/// 引入此标志把两种语义显式分离。
|
||||
count_all: bool,
|
||||
}
|
||||
|
||||
impl RateLimiter {
|
||||
/// 构造通用限流器:仅在鉴权失败(400/401/403)时计数。
|
||||
pub fn new(max_failures: usize, window: Duration) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(HashMap::new())),
|
||||
window,
|
||||
max_failures,
|
||||
count_all: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 构造「全量计数」限流器:对匹配路径的所有响应(无论成败)计数。
|
||||
/// 用于注册端点专用节流。
|
||||
pub fn new_count_all(max_failures: usize, window: Duration) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(HashMap::new())),
|
||||
window,
|
||||
max_failures,
|
||||
count_all: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查该 IP 是否已被限流(窗口内失败次数超阈值)。不修改计数。
|
||||
pub(crate) fn is_rate_limited(&self, ip: IpAddr) -> bool {
|
||||
let now = Instant::now();
|
||||
let mut map = match self.inner.lock() {
|
||||
Ok(g) => g,
|
||||
Err(e) => e.into_inner(), // poisoned:仍尽力返回判断,避免鉴权因锁中毒全部放行
|
||||
};
|
||||
if let Some(queue) = map.get_mut(&ip) {
|
||||
// 清理过期时间戳
|
||||
while let Some(front) = queue.front() {
|
||||
if now.duration_since(*front) > self.window {
|
||||
queue.pop_front();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if queue.is_empty() {
|
||||
map.remove(&ip);
|
||||
return false;
|
||||
}
|
||||
return queue.len() >= self.max_failures;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// 记录一次失败(追加时间戳)。
|
||||
pub(crate) fn record_failure(&self, ip: IpAddr) {
|
||||
let now = Instant::now();
|
||||
let mut map = match self.inner.lock() {
|
||||
Ok(g) => g,
|
||||
Err(e) => e.into_inner(),
|
||||
};
|
||||
let queue = map.entry(ip).or_default();
|
||||
queue.push_back(now);
|
||||
// 顺带清理,防止队列无限增长
|
||||
while let Some(front) = queue.front() {
|
||||
if now.duration_since(*front) > self.window {
|
||||
queue.pop_front();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if queue.is_empty() {
|
||||
map.remove(&ip);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 判断 IP 是否为本地环回或私有网段 IP。
|
||||
fn is_private_or_loopback_ip(ip: IpAddr) -> bool {
|
||||
match ip {
|
||||
IpAddr::V4(v4) => v4.is_loopback() || v4.is_private(),
|
||||
IpAddr::V6(v6) => v6.is_loopback(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从请求中提取客户端 IP。
|
||||
/// 仅当底层连接 (ConnectInfo) 为本地环回或私有网段时才信任反向代理传递的 X-Forwarded-For / X-Real-IP。
|
||||
fn extract_client_ip(req: &Request<axum::body::Body>) -> Option<IpAddr> {
|
||||
let direct_ip = req
|
||||
.extensions()
|
||||
.get::<ConnectInfo<SocketAddr>>()
|
||||
.map(|ci| ci.0.ip());
|
||||
|
||||
// 如果有直连 IP 且不是私有/环回地址,说明未经过可信反代,直接返回直连 IP 拒绝盲信 X-Forwarded-For
|
||||
if let Some(ip) = direct_ip {
|
||||
if !is_private_or_loopback_ip(ip) {
|
||||
return Some(ip);
|
||||
}
|
||||
}
|
||||
|
||||
// 只有处于本地/私有网络反代之后时,才尝试提取 X-Forwarded-For
|
||||
if let Some(xff) = req
|
||||
.headers()
|
||||
.get("x-forwarded-for")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
if let Some(first) = xff.split(',').map(|s| s.trim()).next() {
|
||||
if !first.is_empty() {
|
||||
if let Ok(ip) = first.parse::<IpAddr>() {
|
||||
return Some(ip);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// 回退:X-Real-IP
|
||||
if let Some(xri) = req.headers().get("x-real-ip").and_then(|v| v.to_str().ok()) {
|
||||
if let Ok(ip) = xri.parse::<IpAddr>() {
|
||||
return Some(ip);
|
||||
}
|
||||
}
|
||||
// 回退:直连 IP
|
||||
direct_ip
|
||||
}
|
||||
|
||||
/// 限流中间件:在鉴权之前检查该 IP 是否已被限流。
|
||||
///
|
||||
/// 放在 auth_middleware **之前**(外层):被限流的 IP 直接 429,不进鉴权逻辑。
|
||||
/// 是否记入失败窗口,由 auth_middleware 的结果决定——为此 auth 中间件会把 401 的 IP
|
||||
/// 通过本 limiter 记录。但为避免跨中间件传参的复杂性,这里采用「先放行让 auth 判定,
|
||||
/// 若返回 401 再记录」的方式:见下方包装函数 `rate_limit_with_auth`。
|
||||
pub async fn rate_limit_middleware(
|
||||
State(limiter): State<RateLimiter>,
|
||||
req: Request<axum::body::Body>,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
let ip = extract_client_ip(&req).unwrap_or(IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED));
|
||||
let is_register = req.uri().path().ends_with("/node/register");
|
||||
|
||||
if limiter.is_rate_limited(ip) {
|
||||
warn!("客户端 IP {} 鉴权失败次数过多,已限流(429)", ip);
|
||||
return (
|
||||
axum::http::StatusCode::TOO_MANY_REQUESTS,
|
||||
"鉴权失败次数过多,请稍后重试",
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
let resp = next.run(req).await;
|
||||
|
||||
// 计入速率窗口的条件:
|
||||
// - 鉴权失败(401/403/400):通用与专用限流器都计;
|
||||
// - 或 limiter 配置为 count_all 且请求落在专用节流路径(如 /node/register):
|
||||
// 这种情况下成功响应也计,作为对注册接口本身的独立节流(防恶意频繁注册)。
|
||||
// 通用 API 限流器(count_all=false)不会因 is_register 把成功请求计入,
|
||||
// 避免了「成功注册连锁锁出整个 /api/*」的历史缺陷。
|
||||
let status = resp.status().as_u16();
|
||||
let auth_failed = status == 401 || status == 403 || status == 400;
|
||||
if auth_failed || (limiter.count_all && is_register) {
|
||||
limiter.record_failure(ip);
|
||||
}
|
||||
|
||||
resp
|
||||
}
|
||||
@@ -2,7 +2,7 @@ use super::AppState;
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::{Path as AxumPath, State},
|
||||
http::{header, StatusCode},
|
||||
http::header,
|
||||
response::Response,
|
||||
};
|
||||
use tokio::fs::File;
|
||||
@@ -12,13 +12,20 @@ use tracing::warn;
|
||||
pub async fn download_seed(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> Response {
|
||||
if name.is_empty() || name.starts_with('.') || !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@') {
|
||||
warn!("拒绝可能包含路径穿越或特别注入序列的非法种子下载请求: {}", name);
|
||||
return Response::builder()
|
||||
.status(StatusCode::BAD_REQUEST)
|
||||
.body(Body::from("非法的种子名称参数"))
|
||||
.unwrap();
|
||||
) -> Result<Response, crate::api::AppError> {
|
||||
if name.is_empty()
|
||||
|| name.starts_with('.')
|
||||
|| !name.chars().all(|c| {
|
||||
c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@'
|
||||
})
|
||||
{
|
||||
warn!(
|
||||
"拒绝可能包含路径穿越或特别注入序列的非法种子下载请求: {}",
|
||||
name
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的种子名称参数".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let seed_file_path = std::path::Path::new(&state.results_dir)
|
||||
@@ -27,32 +34,27 @@ pub async fn download_seed(
|
||||
|
||||
if !seed_file_path.is_file() {
|
||||
warn!("客户端请求的种子文件不存在: {}", seed_file_path.display());
|
||||
return Response::builder()
|
||||
.status(StatusCode::NOT_FOUND)
|
||||
.body(Body::from("请求的种子文件不存在"))
|
||||
.unwrap();
|
||||
return Err(crate::api::AppError::NotFound(
|
||||
"请求的种子文件不存在".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let file = match File::open(&seed_file_path).await {
|
||||
Ok(file) => file,
|
||||
Err(_) => {
|
||||
return Response::builder()
|
||||
.status(StatusCode::INTERNAL_SERVER_ERROR)
|
||||
.body(Body::from("无法打开种子文件"))
|
||||
.unwrap();
|
||||
Err(e) => {
|
||||
return Err(crate::api::AppError::Internal(e.into()));
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
let stream = ReaderStream::new(file);
|
||||
let body = Body::from_stream(stream);
|
||||
|
||||
Response::builder()
|
||||
Ok(Response::builder()
|
||||
.header(header::CONTENT_TYPE, "application/octet-stream")
|
||||
.header(
|
||||
header::CONTENT_DISPOSITION,
|
||||
format!("attachment; filename=\"{}.7\"", name),
|
||||
)
|
||||
.body(body)
|
||||
.unwrap()
|
||||
.unwrap())
|
||||
}
|
||||
|
||||
@@ -2,21 +2,37 @@ use super::AppState;
|
||||
use axum::{extract::State, response::IntoResponse, Json};
|
||||
use serde_json::json;
|
||||
|
||||
pub async fn get_status(State(state): State<AppState>) -> impl IntoResponse {
|
||||
/// 轻量健康检查端点(不走鉴权)。
|
||||
///
|
||||
/// 供 docker healthcheck、负载均衡、外部监控探测。刻意只返回固定 ok,
|
||||
/// 不触碰数据库或调度器,避免健康检查本身拖累系统或因 DB 瞬时锁导致误判不健康。
|
||||
pub async fn healthz() -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
Ok(Json(json!({ "status": "ok" })))
|
||||
}
|
||||
|
||||
pub async fn get_status(
|
||||
State(state): State<AppState>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
let nodes = state.db.get_active_nodes().await.unwrap_or_default();
|
||||
let total_active_slots: i32 = nodes.iter().map(|n| n.active_slots).sum();
|
||||
let total_max_slots: i32 = nodes.iter().map(|n| n.max_slots).sum();
|
||||
|
||||
let grid_stats = state.db.get_grid_summary_stats().await.unwrap_or(serde_json::json!({
|
||||
"total": 0, "pending": 0, "running": 0, "converged": 0, "failed": 0
|
||||
}));
|
||||
// dashboard 全局概览:聚合全部工作流的 grid_points(多工作流分区后仍提供全局合计)。
|
||||
// 若需单工作流进度,可扩展为按 workflow 查询参数分别聚合。
|
||||
let grid_stats = state
|
||||
.db
|
||||
.get_grid_summary_stats(None)
|
||||
.await
|
||||
.unwrap_or(serde_json::json!({
|
||||
"total": 0, "pending": 0, "running": 0, "converged": 0, "failed": 0
|
||||
}));
|
||||
|
||||
Json(json!({
|
||||
Ok(Json(json!({
|
||||
"status": "online",
|
||||
"nodes_online": nodes.len(),
|
||||
"total_active_slots": total_active_slots,
|
||||
"total_max_slots": total_max_slots,
|
||||
"nodes": nodes,
|
||||
"grid_stats": grid_stats,
|
||||
}))
|
||||
})))
|
||||
}
|
||||
|
||||
+125
-54
@@ -1,6 +1,6 @@
|
||||
use super::AppState;
|
||||
use super::{AppState, AuthenticatedNode};
|
||||
use axum::{
|
||||
extract::{Multipart, State},
|
||||
extract::{Extension, Multipart, State},
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
@@ -12,27 +12,38 @@ use tracing::{info, warn};
|
||||
|
||||
use axum::http::StatusCode;
|
||||
|
||||
pub async fn claim_task(State(state): State<AppState>) -> impl IntoResponse {
|
||||
match state.queue.pop_task().await {
|
||||
pub async fn claim_task(
|
||||
State(state): State<AppState>,
|
||||
Extension(auth_node): Extension<AuthenticatedNode>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// 领用时记录任务归属:pop_task 写入 claimed_by_node_id,
|
||||
// report 阶段据此校验「上报者确为领用者」,杜绝跨节点伪造结果。
|
||||
match state.queue.pop_task(&auth_node.node_id).await {
|
||||
Ok(Some(task)) => {
|
||||
if let Err(e) = state.db.mark_grid_point_running(&task.point_name).await {
|
||||
// 多工作流分区:mark_grid_point_running 须带 workflow_name,避免按 name 全局更新
|
||||
// 误改其他工作流的同名点。TaskSpec.workflow_name 在调度时已绑定。
|
||||
let wf = task.workflow_name.as_deref().unwrap_or("");
|
||||
if let Err(e) = state.db.mark_grid_point_running(&task.point_name, wf).await {
|
||||
warn!("领用任务 {} 后同步变更为 running 状态遇到异常: {}. 后置 stale 定时自取检索引索将介入修复维护", task.task_id, e);
|
||||
}
|
||||
(StatusCode::OK, Json(json!({"status": "ok", "task": task}))).into_response()
|
||||
Ok((StatusCode::OK, Json(json!({"status": "ok", "task": task}))))
|
||||
}
|
||||
Ok(None) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({"status": "empty", "task": null})),
|
||||
)),
|
||||
Err(e) => {
|
||||
tracing::error!("领用任务数据库异常: {}", e);
|
||||
Err(crate::api::AppError::Internal(e))
|
||||
}
|
||||
Ok(None) => (StatusCode::OK, Json(json!({"status": "empty", "task": null}))).into_response(),
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({"status": "error", "message": format!("领用任务失败: {}", e)})),
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn report_task(
|
||||
State(state): State<AppState>,
|
||||
Extension(auth_node): Extension<AuthenticatedNode>,
|
||||
mut multipart: Multipart,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
let mut report_json: Option<TaskReport> = None;
|
||||
let mut seed_file_data: Option<Vec<u8>> = None;
|
||||
|
||||
@@ -51,53 +62,98 @@ pub async fn report_task(
|
||||
}
|
||||
}
|
||||
|
||||
let report = match report_json {
|
||||
let mut report = match report_json {
|
||||
Some(r) => r,
|
||||
None => {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({"status": "error", "message": "请求中缺少 report 字段"})),
|
||||
)
|
||||
.into_response();
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"请求中缺少 report 字段".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// ── 任务归属校验(S1 核心,防跨节点伪造结果投毒)──
|
||||
// 1. 该 task_id 必须由当前鉴权 node 领用(claim 时记录的 claimed_by_node_id 匹配)。
|
||||
// 2. 上报的 point_name 必须与该 task 绑定的 point_name 一致(防跨点上报)。
|
||||
// 3. 忽略 body 里声称的 node_id,统一以鉴权 node_id 写库(修复审计归因断裂)。
|
||||
// 4. 取 task 绑定的 workflow_name,用于定向更新该工作流的 grid_points(多工作流分区)。
|
||||
let (claimed_point, claimed_workflow) = match state
|
||||
.queue
|
||||
.verify_task_claim(&report.task_id.to_string(), &auth_node.node_id)
|
||||
.await
|
||||
{
|
||||
Ok(Some((p, w))) => (p, w),
|
||||
Ok(None) => {
|
||||
warn!(
|
||||
"任务归属校验失败:node={} 上报 task_id={} 但未领用或已被清理",
|
||||
auth_node.node_id, report.task_id
|
||||
);
|
||||
return Err(crate::api::AppError::Forbidden(
|
||||
"任务未由本节点领用或已上报过".to_string(),
|
||||
));
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("校验任务归属数据库异常: {}", e);
|
||||
return Err(crate::api::AppError::Internal(e));
|
||||
}
|
||||
};
|
||||
if claimed_point != report.point_name {
|
||||
warn!(
|
||||
"任务点名校验失败:task_id={} 领用 point={} 但上报 point={}",
|
||||
report.task_id, claimed_point, report.point_name
|
||||
);
|
||||
return Err(crate::api::AppError::Forbidden(
|
||||
"上报的网格点与领用任务不匹配".to_string(),
|
||||
));
|
||||
}
|
||||
// 统一以鉴权 node_id 覆盖 body 里的 node_id,保证归因可信
|
||||
report.node_id = auth_node.node_id.clone();
|
||||
// workflow_name 以领用记录为准(claim 时从 TaskSpec 落库),body 无权声称。
|
||||
let workflow_name = claimed_workflow.unwrap_or_default();
|
||||
|
||||
let name = report.point_name.clone();
|
||||
|
||||
if name.is_empty() || name.starts_with('.') || !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@') {
|
||||
warn!("拒绝可能包含路径穿越或特殊非常规编码号攻击的网格点名称请求: {}", name);
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({"status": "error", "message": "非法的网格点名称参数"})),
|
||||
)
|
||||
.into_response();
|
||||
if name.is_empty()
|
||||
|| name.starts_with('.')
|
||||
|| !name.chars().all(|c| {
|
||||
c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@'
|
||||
})
|
||||
{
|
||||
warn!(
|
||||
"拒绝可能包含路径穿越或特殊非常规编码号攻击的网格点名称请求: {}",
|
||||
name
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的网格点名称参数".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let params = match extract_params(&report) {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
warn!("网格点 {} 汇报数据解析失败: 无法解析 params 或 summary_json", name);
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({"status": "error", "message": "无法解析 params 或 summary_json"})),
|
||||
)
|
||||
.into_response();
|
||||
warn!(
|
||||
"网格点 {} 汇报数据解析失败: 无法解析 params 或 summary_json",
|
||||
name
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"无法解析 params 或 summary_json".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// Record in DB
|
||||
if let Err(e) = state.db.record_task_report(&report).await {
|
||||
warn!("记录网格点 {} 任务结果到数据库失败: {}", name, e);
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({"status": "error", "message": format!("记录数据库失败: {}", e)})),
|
||||
)
|
||||
.into_response();
|
||||
// Record in DB(带 workflow_name 定向更新该工作流的 grid_points)
|
||||
if let Err(e) = state.db.record_task_report(&report, &workflow_name).await {
|
||||
// DB 错误细节进日志,对客户端只返回通用消息(避免泄露表结构/内部错误给未授权方)
|
||||
tracing::error!("记录网格点 {} 任务结果到数据库失败: {}", name, e);
|
||||
return Err(crate::api::AppError::Internal(e));
|
||||
}
|
||||
|
||||
// Clean up task from task_queue table to prevent queue DB bloat
|
||||
if let Err(e) = state.queue.remove_task(&report.task_id.to_string()).await {
|
||||
tracing::warn!("从任务队列中清理已上报任务记录 {} 失败: {}", report.task_id, e);
|
||||
tracing::warn!(
|
||||
"从任务队列中清理已上报任务记录 {} 失败: {}",
|
||||
report.task_id,
|
||||
e
|
||||
);
|
||||
}
|
||||
|
||||
// 采用原子写入模式保持 conv.json 与核心二进制数据完整落地后才揭晓真实文件名
|
||||
@@ -112,30 +168,46 @@ pub async fn report_task(
|
||||
// Save seed file .7 using atomic temporary writing strategy
|
||||
if report.converged && !report.atmosphere_has_nan {
|
||||
if let Some(bytes) = seed_file_data {
|
||||
let seed_tmp = model_dir.join(format!("{}.7.{}.tmp", name, uuid::Uuid::new_v4().simple()));
|
||||
let seed_tmp =
|
||||
model_dir.join(format!("{}.7.{}.tmp", name, uuid::Uuid::new_v4().simple()));
|
||||
let seed_path = model_dir.join(format!("{}.7", name));
|
||||
if fs::write(&seed_tmp, bytes).await.is_ok() {
|
||||
if fs::rename(&seed_tmp, &seed_path).await.is_ok() {
|
||||
info!("成功保持原子写入落地并保存网格点 {} 的收敛种子文件: {}", name, seed_path.display());
|
||||
let _ = state
|
||||
.db
|
||||
.insert_seed(¶ms, &seed_path.to_string_lossy())
|
||||
.await;
|
||||
}
|
||||
if fs::write(&seed_tmp, bytes).await.is_ok()
|
||||
&& fs::rename(&seed_tmp, &seed_path).await.is_ok()
|
||||
{
|
||||
info!(
|
||||
"成功保持原子写入落地并保存网格点 {} 的收敛种子文件: {}",
|
||||
name,
|
||||
seed_path.display()
|
||||
);
|
||||
let _ = state
|
||||
.db
|
||||
.insert_seed(¶ms, &seed_path.to_string_lossy())
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if report.status == TaskStatus::Failed || report.status == TaskStatus::Timeout || report.atmosphere_has_nan {
|
||||
if !report.converged
|
||||
|| report.atmosphere_has_nan
|
||||
|| report.status == TaskStatus::Failed
|
||||
|| report.status == TaskStatus::Timeout
|
||||
{
|
||||
// Task did not succeed -> check if seed_step fallback should be triggered
|
||||
info!("网格点 {} 计算未成功完成,检查种子回退机制...", name);
|
||||
if let Err(e) = state.scheduler.trigger_seed_step_fallback(¶ms).await {
|
||||
if let Err(e) = state
|
||||
.scheduler
|
||||
.trigger_seed_step_fallback(¶ms, &workflow_name)
|
||||
.await
|
||||
{
|
||||
warn!("网格点 {} 触发种子回退机制失败: {}", name, e);
|
||||
}
|
||||
}
|
||||
|
||||
(StatusCode::OK, Json(json!({"status": "ok", "message": "上报成功"}))).into_response()
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({"status": "ok", "message": "上报成功"})),
|
||||
))
|
||||
}
|
||||
|
||||
fn extract_params(report: &TaskReport) -> Option<GridPointParams> {
|
||||
@@ -146,4 +218,3 @@ fn extract_params(report: &TaskReport) -> Option<GridPointParams> {
|
||||
.ok()
|
||||
.map(|summary| summary.params)
|
||||
}
|
||||
|
||||
|
||||
+141
-138
@@ -23,162 +23,166 @@ pub struct ApiResponse<T> {
|
||||
pub data: Option<T>,
|
||||
}
|
||||
|
||||
pub async fn list_workflows(State(state): State<AppState>) -> impl IntoResponse {
|
||||
pub async fn list_workflows(
|
||||
State(state): State<AppState>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
match state.db.list_workflows().await {
|
||||
Ok(list) => (StatusCode::OK, Json(ApiResponse { success: true, message: "成功获取工作流列表".to_string(), data: Some(list) })),
|
||||
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(ApiResponse { success: false, message: format!("获取工作流列表失败: {}", e), data: None })),
|
||||
Ok(list) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse {
|
||||
success: true,
|
||||
message: "成功获取工作流列表".to_string(),
|
||||
data: Some(list),
|
||||
}),
|
||||
)),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
match state.db.get_workflow(&name).await {
|
||||
Ok(Some(item)) => (StatusCode::OK, Json(ApiResponse { success: true, message: "成功获取工作流详情".to_string(), data: Some(item) })),
|
||||
Ok(None) => (StatusCode::NOT_FOUND, Json(ApiResponse { success: false, message: format!("工作流 '{}' 未找到", name), data: None })),
|
||||
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(ApiResponse { success: false, message: format!("获取工作流详情失败: {}", e), data: None })),
|
||||
Ok(Some(item)) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse {
|
||||
success: true,
|
||||
message: "成功获取工作流详情".to_string(),
|
||||
data: Some(item),
|
||||
}),
|
||||
)),
|
||||
Ok(None) => Err(crate::api::AppError::NotFound(format!(
|
||||
"工作流 '{}' 未找到",
|
||||
name
|
||||
))),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 工作流名称白名单:仅允许字母、数字、点、下划线、连字符,长度 1-64。
|
||||
/// 与 report_task/download_seed 的网格点名校验口径保持一致,从源头阻止
|
||||
/// 名称携带 HTML/JS 特殊字符进入 Dashboard 渲染(存储型 XSS 根因之一)。
|
||||
fn is_valid_workflow_name(name: &str) -> bool {
|
||||
!name.is_empty()
|
||||
&& name.len() <= 64
|
||||
&& name
|
||||
.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-')
|
||||
}
|
||||
|
||||
pub async fn save_workflow(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<CreateWorkflowRequest>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// 名称白名单校验(优先于 YAML 校验,拒绝携带特殊字符的名称)
|
||||
if !is_valid_workflow_name(&req.name) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"工作流名称仅允许字母、数字、点(.)、下划线(_)、连字符(-),长度 1-64".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Validate YAML config string
|
||||
if let Err(e) = serde_yaml::from_str::<GridConfig>(&req.config_yaml) {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("无效的 YAML 配置: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(format!(
|
||||
"无效的 YAML 配置: {}",
|
||||
e
|
||||
)));
|
||||
}
|
||||
|
||||
// 检查被编辑的工作流是否正处于激活运行中
|
||||
if let Ok(Some(existing)) = state.db.get_workflow(&req.name).await {
|
||||
if existing.status == "running" || existing.status == "initializing" {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 正处在运行或初始加载流程中,严禁原地覆写参数重设至 IDLE;如待变更参数请先调 API 显式触发停止后再保存", req.name),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
format!("工作流 '{}' 正处在运行或初始加载流程中,严禁原地覆写参数重设至 IDLE;如待变更参数请先调 API 显式触发停止后再保存", req.name)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
match state.db.upsert_workflow(&req.name, req.description.as_deref(), &req.config_yaml, "idle").await {
|
||||
match state
|
||||
.db
|
||||
.upsert_workflow(
|
||||
&req.name,
|
||||
req.description.as_deref(),
|
||||
&req.config_yaml,
|
||||
"idle",
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
info!("成功注册/更新工作流配置: {}", req.name);
|
||||
(
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 保存成功", req.name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
))
|
||||
}
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("保存工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn delete_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// 拦截正在运行或初始加载中的工作流删除请求
|
||||
if let Ok(Some(existing)) = state.db.get_workflow(&name).await {
|
||||
if existing.status == "running" || existing.status == "initializing" {
|
||||
return Err(crate::api::AppError::BadRequest(format!(
|
||||
"工作流 '{}' 当前处于 '{}' 状态,无法直接删除。请先显式暂停/停止该工作流。",
|
||||
name, existing.status
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
let _ = state.queue.clear_queue_by_workflow(&name).await;
|
||||
match state.db.delete_workflow(&name).await {
|
||||
Ok(_) => (
|
||||
Ok(_) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 已删除", name),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("删除工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
)),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn start_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
let item = match state.db.get_workflow(&name).await {
|
||||
Ok(Some(item)) => item,
|
||||
Ok(None) => {
|
||||
return (
|
||||
StatusCode::NOT_FOUND,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 未找到", name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
}
|
||||
Err(e) => {
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("获取工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
return Err(crate::api::AppError::NotFound(format!(
|
||||
"工作流 '{}' 未找到",
|
||||
name
|
||||
)))
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
|
||||
if item.status == "running" || item.status == "initializing" {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 已处在初始建立状态中或者已处于运行状态,无需且不允许进行并行重置启动", name),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(format!(
|
||||
"工作流 '{}' 已处在初始建立状态中或者已处于运行状态,无需且不允许进行并行重置启动",
|
||||
name
|
||||
)));
|
||||
}
|
||||
|
||||
// 通过原子性抢占更新将状态切换为 initializing,拦截同名流上的多并发调用导致的双重加载破坏性竞态
|
||||
match state.db.transition_workflow_to_initializing(&name).await {
|
||||
Ok(false) => {
|
||||
return (
|
||||
StatusCode::CONFLICT,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 初始化抢占挂起异常,表明已在另一会话上下文中顺利推入启动通道", name),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("原子化抢占和迁移工作流状态发生异常: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
return Err(crate::api::AppError::Conflict(format!(
|
||||
"工作流 '{}' 初始化抢占挂起异常,表明已在另一会话上下文中顺利推入启动通道",
|
||||
name
|
||||
)));
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
Ok(true) => {}
|
||||
}
|
||||
|
||||
@@ -186,70 +190,69 @@ pub async fn start_workflow(
|
||||
Ok(cfg) => cfg,
|
||||
Err(e) => {
|
||||
let _ = state.db.update_workflow_status(&name, "idle").await;
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("解析工作流 YAML 发生语法或参数解析异常: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(format!(
|
||||
"解析工作流 YAML 发生语法或参数解析异常: {}",
|
||||
e
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
info!("成功占据独享启动权,开始启动工作流 '{}',系统进行 64/32 维深度平展开网格结构计算化推列并推送队列...", name);
|
||||
match state.scheduler.initialize_grid(&grid_cfg).await {
|
||||
Ok(_) => {
|
||||
let _ = state.db.update_workflow_status(&name, "running").await;
|
||||
let _ = state.scheduler.schedule_pending_tasks().await;
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 建立与挂载成功并已接续排班", name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
info!("成功占据独享启动权,开始异步启动工作流 '{}',系统将在后台进行 64/32 维深度平展开网格结构计算化推列并推送队列...", name);
|
||||
|
||||
let bg_state = state.clone();
|
||||
let bg_name = name.clone();
|
||||
let bg_grid_cfg = grid_cfg;
|
||||
|
||||
tokio::spawn(async move {
|
||||
match bg_state
|
||||
.scheduler
|
||||
.initialize_grid(&bg_grid_cfg, &bg_name)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
let _ = bg_state
|
||||
.db
|
||||
.update_workflow_status(&bg_name, "running")
|
||||
.await;
|
||||
let _ = bg_state.scheduler.schedule_pending_tasks().await;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("工作流 {} 网格初始化中途失败,已回退为 idle;已写入的点保留,重新启动会幂等补齐: {}", bg_name, e);
|
||||
let _ = bg_state.db.update_workflow_status(&bg_name, "idle").await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = state.db.update_workflow_status(&name, "idle").await;
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("展开与挂载初始化任务点到系统队列失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 已进入后台异步建立与挂载流程", name),
|
||||
data: None,
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
pub async fn stop_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
match state.db.update_workflow_status(&name, "paused").await {
|
||||
Ok(_) => {
|
||||
let _ = state.queue.clear_queue().await;
|
||||
let _ = state.db.reset_queued_grid_points_to_pending().await;
|
||||
(
|
||||
// 多工作流分区:清理与重置都限定在本工作流内,避免误伤其他并发运行的工作流。
|
||||
// - clear_queue_by_workflow:只删本工作流的排队任务。
|
||||
// - reset_queued_grid_points_to_pending(&name):只把本工作流的 queued 点打回 pending。
|
||||
let _ = state.queue.clear_queue_by_workflow(&name).await;
|
||||
let _ = state.db.reset_queued_grid_points_to_pending(&name).await;
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 已暂停,排队任务已暂停调度", name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
))
|
||||
}
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("暂停工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
use axum::http::HeaderValue;
|
||||
use tower_http::cors::{AllowOrigin, CorsLayer};
|
||||
use tracing::info;
|
||||
|
||||
/// 构建 CORS 中间件层。
|
||||
///
|
||||
/// 严格安全策略:仅允许**同源**(Origin 匹配请求头的 Host)或**本地 Origin**(localhost / 127.0.0.1 / [::1])。
|
||||
pub fn build_cors_layer() -> CorsLayer {
|
||||
info!("CORS 策略:仅允许同源或本地 Origin(localhost / 127.0.0.1 / [::1])");
|
||||
|
||||
CorsLayer::new()
|
||||
.allow_origin(AllowOrigin::predicate(
|
||||
|origin: &HeaderValue, head: &axum::http::request::Parts| {
|
||||
let Ok(origin_str) = origin.to_str() else {
|
||||
return false;
|
||||
};
|
||||
let Ok(uri) = origin_str.parse::<axum::http::Uri>() else {
|
||||
return false;
|
||||
};
|
||||
let Some(host) = uri.host() else {
|
||||
return false;
|
||||
};
|
||||
let clean_host = host.trim_start_matches('[').trim_end_matches(']');
|
||||
|
||||
// 1. 本地来源 (localhost / 127.0.0.1 / [::1])
|
||||
if clean_host == "localhost"
|
||||
|| clean_host == "127.0.0.1"
|
||||
|| clean_host == "::1"
|
||||
|| clean_host.starts_with("127.")
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
// 2. 同源来源 (Origin 匹配请求头的 Host)
|
||||
if let Some(host_header) = head.headers.get(axum::http::header::HOST) {
|
||||
if let Ok(host_str) = host_header.to_str() {
|
||||
if let Some(authority) = uri.authority() {
|
||||
if authority.as_str().eq_ignore_ascii_case(host_str) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
},
|
||||
))
|
||||
.allow_methods([
|
||||
axum::http::Method::GET,
|
||||
axum::http::Method::POST,
|
||||
axum::http::Method::PUT,
|
||||
axum::http::Method::DELETE,
|
||||
])
|
||||
.allow_headers([
|
||||
axum::http::header::AUTHORIZATION,
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
])
|
||||
}
|
||||
+1354
-106
File diff suppressed because it is too large
Load Diff
@@ -1,4 +1,4 @@
|
||||
pub mod api;
|
||||
pub mod cors;
|
||||
pub mod db;
|
||||
pub mod scheduler;
|
||||
|
||||
|
||||
+280
-43
@@ -3,8 +3,9 @@ use server::api::{self, AppState};
|
||||
use server::db::Database;
|
||||
use server::scheduler::GridScheduler;
|
||||
|
||||
|
||||
use axum::{
|
||||
extract::DefaultBodyLimit,
|
||||
http::HeaderValue,
|
||||
routing::{get, post},
|
||||
Router,
|
||||
};
|
||||
@@ -16,12 +17,15 @@ use std::net::SocketAddr;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use tokio::time::{sleep, Duration};
|
||||
use tower_http::cors::CorsLayer;
|
||||
use tower_http::services::{ServeDir, ServeFile};
|
||||
use tracing::info;
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "server", version = "0.1.0", about = "Distributed Computing TLUSTY/SYNSPEC (DCTS) Server")]
|
||||
#[command(
|
||||
name = "server",
|
||||
version = "0.1.0",
|
||||
about = "Distributed Computing TLUSTY/SYNSPEC (DCTS) Server"
|
||||
)]
|
||||
struct CliArgs {
|
||||
/// Optional path to workflow configuration YAML file to auto-register on startup
|
||||
#[arg(short = 'w', long = "workflow")]
|
||||
@@ -61,12 +65,15 @@ async fn main() -> Result<()> {
|
||||
let default_wf_path = Path::new(&server_cfg.grid_config);
|
||||
if default_wf_path.is_file() {
|
||||
if let Ok(yaml_content) = std::fs::read_to_string(default_wf_path) {
|
||||
if let Err(e) = db.upsert_workflow(
|
||||
"sdB_cno",
|
||||
Some("sdB CNO 6D Stellar Atmosphere Grid"),
|
||||
&yaml_content,
|
||||
"idle",
|
||||
).await {
|
||||
if let Err(e) = db
|
||||
.upsert_workflow(
|
||||
"sdB_cno",
|
||||
Some("sdB CNO 6D Stellar Atmosphere Grid"),
|
||||
&yaml_content,
|
||||
"idle",
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("预注册默认工作流失败: {}", e);
|
||||
} else {
|
||||
info!("已在数据库中成功预注册默认工作流 'sdB_cno'");
|
||||
@@ -74,15 +81,38 @@ async fn main() -> Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
// 弱口令凭据安全警告检测
|
||||
let is_weak_token = |t: Option<&str>| -> bool {
|
||||
match t {
|
||||
Some(s) => {
|
||||
s.len() < 12 || s == "fmqi123" || s == "admin" || s == "123456" || s == "secret"
|
||||
}
|
||||
None => false,
|
||||
}
|
||||
};
|
||||
if is_weak_token(server_cfg.auth_token.as_deref())
|
||||
|| is_weak_token(server_cfg.admin_token.as_deref())
|
||||
{
|
||||
tracing::warn!("⚠️ 检测到系统当前正在使用弱口令凭据或默认 Token!建议生产环境在 .env 中配置使用 openssl rand -hex 32 生成的高强度 Token!");
|
||||
}
|
||||
|
||||
let rate_limiter = api::rate_limit::RateLimiter::new(5, std::time::Duration::from_secs(300));
|
||||
|
||||
let state = AppState {
|
||||
db,
|
||||
queue: queue.clone(),
|
||||
scheduler: scheduler.clone(),
|
||||
results_dir: server_cfg.results_dir,
|
||||
results_dir: server_cfg.results_dir.clone(),
|
||||
rate_limiter,
|
||||
auth_token: server_cfg.auth_token.clone(),
|
||||
admin_token: server_cfg.admin_token.clone(),
|
||||
auth_disabled: server_cfg.auth_disabled,
|
||||
admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||||
std::collections::HashMap::new(),
|
||||
)),
|
||||
};
|
||||
|
||||
// Background loop for stale task requeueing, offline node detection, and scheduler checking
|
||||
// Background maintenance & scheduling with Exponential Backoff
|
||||
let bg_db = state.db.clone();
|
||||
let bg_queue = queue.clone();
|
||||
let bg_scheduler = scheduler.clone();
|
||||
@@ -90,58 +120,226 @@ async fn main() -> Result<()> {
|
||||
let node_stale_sec = server_cfg.node_stale_sec;
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut fail_count: u32 = 0;
|
||||
let mut first_run = true;
|
||||
loop {
|
||||
sleep(Duration::from_secs(30)).await;
|
||||
if let Ok(requeued_points) = bg_queue.requeue_stale_tasks(stale_sec).await {
|
||||
if !requeued_points.is_empty() {
|
||||
info!("重新将 {} 个超时/掉线任务放回待计算队列", requeued_points.len());
|
||||
let _ = bg_db.reset_specific_grid_points_to_pending(&requeued_points).await;
|
||||
}
|
||||
if first_run {
|
||||
first_run = false;
|
||||
} else {
|
||||
let base_delay = 30u64;
|
||||
let current_delay = if fail_count == 0 {
|
||||
base_delay
|
||||
} else {
|
||||
(base_delay * (1u64 << fail_count.min(4))).min(300)
|
||||
};
|
||||
sleep(Duration::from_secs(current_delay)).await;
|
||||
}
|
||||
if let Ok(offline) = bg_db.mark_stale_nodes_offline(node_stale_sec).await {
|
||||
if offline > 0 {
|
||||
info!("已标记 {} 个心跳超时的计算节点为离线状态", offline);
|
||||
|
||||
let bg_db_clone = bg_db.clone();
|
||||
let bg_queue_clone = bg_queue.clone();
|
||||
let bg_scheduler_clone = bg_scheduler.clone();
|
||||
|
||||
let join_handle = tokio::spawn(async move {
|
||||
let mut has_error = false;
|
||||
match bg_queue_clone.requeue_stale_tasks(stale_sec).await {
|
||||
Ok(requeued) => {
|
||||
if !requeued.is_empty() {
|
||||
info!("重新将 {} 个超时/掉线任务放回待计算队列", requeued.len());
|
||||
let mut by_wf: std::collections::HashMap<String, Vec<String>> =
|
||||
std::collections::HashMap::new();
|
||||
for (point, wf) in &requeued {
|
||||
by_wf
|
||||
.entry(wf.clone().unwrap_or_default())
|
||||
.or_default()
|
||||
.push(point.clone());
|
||||
}
|
||||
for (wf, points) in by_wf {
|
||||
let _ = bg_db_clone
|
||||
.reset_specific_grid_points_to_pending(&points, &wf)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("重投超时任务失败: {}", e);
|
||||
has_error = true;
|
||||
}
|
||||
}
|
||||
|
||||
match bg_db_clone.mark_stale_nodes_offline(node_stale_sec).await {
|
||||
Ok(offline) => {
|
||||
if offline > 0 {
|
||||
info!("已标记 {} 个心跳超时的计算节点为离线状态", offline);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("标记超时节点离线失败: {}", e);
|
||||
has_error = true;
|
||||
}
|
||||
}
|
||||
|
||||
if let Err(e) = bg_scheduler_clone.schedule_pending_tasks().await {
|
||||
tracing::warn!("后台定时性任务调度检测失败: {}", e);
|
||||
has_error = true;
|
||||
}
|
||||
|
||||
if let Err(e) = bg_db_clone.sync_all_running_workflows_completion().await {
|
||||
tracing::warn!("后台同步已完成工作流状态失败: {}", e);
|
||||
has_error = true;
|
||||
}
|
||||
|
||||
has_error
|
||||
});
|
||||
|
||||
match join_handle.await {
|
||||
Ok(has_error) => {
|
||||
if has_error {
|
||||
fail_count = fail_count.saturating_add(1);
|
||||
} else {
|
||||
fail_count = 0;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("后台维护任务内部发生 Panic: {:?}", e);
|
||||
fail_count = fail_count.saturating_add(1);
|
||||
}
|
||||
}
|
||||
if let Err(e) = bg_scheduler.schedule_pending_tasks().await {
|
||||
tracing::warn!("后台定时性任务调度检测失败: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// 每天自动触发一次数据库备份。
|
||||
// 备份目录跟随 server_cfg.backup_dir(DCTS_BACKUP_DIR,默认 data/backups),
|
||||
// 与 DB_PATH 解耦,避免 DB 卷与备份卷不一致时备份落到未持久化层。
|
||||
// 首次延迟 1 小时,避免频繁重启(如调试阶段)短时间堆积备份文件;backup_database
|
||||
// 自身还带有 7 天保留期清理兜底。
|
||||
let backup_db = state.db.clone();
|
||||
let backup_dir = server_cfg.backup_dir.clone();
|
||||
tokio::spawn(async move {
|
||||
sleep(Duration::from_secs(3600)).await;
|
||||
loop {
|
||||
if let Err(e) = backup_db.backup_database(&backup_dir).await {
|
||||
tracing::warn!("自动备份数据库失败: {}", e);
|
||||
}
|
||||
sleep(Duration::from_secs(24 * 3600)).await;
|
||||
}
|
||||
});
|
||||
|
||||
// 大体积上传端点单独拎出,套用更宽松的 body limit(256MB,覆盖收敛种子 .7 文件量级)
|
||||
// 并限制并发数:每个 report 请求最多 256MB 驻留内存,无并发上限时 N 个请求可耗尽内存。
|
||||
// 限流后超出并发数的请求排队等待(而非直接拒绝),保证正常业务不被误伤。
|
||||
// 其余 API 用 10MB 默认上限,防止大文件内存耗尽 DoS。
|
||||
const REPORT_BODY_LIMIT: usize = 256 * 1024 * 1024;
|
||||
const DEFAULT_BODY_LIMIT: usize = 10 * 1024 * 1024;
|
||||
const REPORT_MAX_CONCURRENCY: usize = 4;
|
||||
|
||||
let report_router = Router::new()
|
||||
.route("/task/report", post(api::task::report_task))
|
||||
.layer(DefaultBodyLimit::max(REPORT_BODY_LIMIT))
|
||||
.layer(tower::ServiceBuilder::new().concurrency_limit(REPORT_MAX_CONCURRENCY));
|
||||
|
||||
// 节点注册接口独立 IP 限流保护(每分钟最多 10 次申请,无论成败都计数,防恶意频繁注册)
|
||||
// 使用 new_count_all:此 limiter 专挂 /node/register,对注册路径的所有响应计入窗口。
|
||||
// 通用 API 限流器(见下方 auth_enabled 分支)用 new 构造(count_all=false),不会因
|
||||
// 成功注册把 IP 锁出整个 /api/*,避免跨端点连锁限流。
|
||||
let register_limiter =
|
||||
api::rate_limit::RateLimiter::new_count_all(10, std::time::Duration::from_secs(60));
|
||||
let register_rate_limit_layer = axum::middleware::from_fn_with_state(
|
||||
register_limiter,
|
||||
api::rate_limit::rate_limit_middleware,
|
||||
);
|
||||
|
||||
let api_router = Router::new()
|
||||
// Auth API
|
||||
.route("/login", post(api::auth::login))
|
||||
.route("/auth/check", get(api::auth::check_auth))
|
||||
// Core Node & Task API
|
||||
.route("/node/register", post(api::node::register_node))
|
||||
.route(
|
||||
"/node/register",
|
||||
post(api::node::register_node).layer(register_rate_limit_layer),
|
||||
)
|
||||
.route("/node/check_status", post(api::node::check_node_status))
|
||||
.route("/node/heartbeat", post(api::node::heartbeat_node))
|
||||
.route("/task/claim", post(api::task::claim_task))
|
||||
.route("/task/report", post(api::task::report_task))
|
||||
.route("/seed/:name", get(api::seed::download_seed))
|
||||
.route("/status", get(api::status::get_status))
|
||||
// Static Data API
|
||||
.route("/data/file/*filename", get(api::data::download_single_data_file))
|
||||
.route(
|
||||
"/data/file/*filename",
|
||||
get(api::data::download_single_data_file),
|
||||
)
|
||||
.route("/data/linelist", get(api::data::download_linelist))
|
||||
// Workflow Management CRUD API
|
||||
.route("/workflows", get(api::workflow::list_workflows).post(api::workflow::save_workflow))
|
||||
.route("/workflows/:name", get(api::workflow::get_workflow).put(api::workflow::save_workflow).delete(api::workflow::delete_workflow))
|
||||
.route("/workflows/:name/start", post(api::workflow::start_workflow))
|
||||
.route("/workflows/:name/stop", post(api::workflow::stop_workflow));
|
||||
.route(
|
||||
"/workflows",
|
||||
get(api::workflow::list_workflows).post(api::workflow::save_workflow),
|
||||
)
|
||||
.route(
|
||||
"/workflows/:name",
|
||||
get(api::workflow::get_workflow)
|
||||
.put(api::workflow::save_workflow)
|
||||
.delete(api::workflow::delete_workflow),
|
||||
)
|
||||
.route(
|
||||
"/workflows/:name/start",
|
||||
post(api::workflow::start_workflow),
|
||||
)
|
||||
.route("/workflows/:name/stop", post(api::workflow::stop_workflow))
|
||||
// Admin Management API(节点凭据查看/审批/吊销/重发,均要求 Admin 角色)
|
||||
.route("/admin/nodes", get(api::admin::list_nodes))
|
||||
.route(
|
||||
"/admin/nodes/:node_id/approve",
|
||||
post(api::admin::approve_node),
|
||||
)
|
||||
.route(
|
||||
"/admin/nodes/:node_id/reject",
|
||||
post(api::admin::reject_node),
|
||||
)
|
||||
.route(
|
||||
"/admin/nodes/:node_id/revoke",
|
||||
post(api::admin::revoke_node),
|
||||
)
|
||||
.route(
|
||||
"/admin/nodes/:node_id/reissue",
|
||||
post(api::admin::reissue_node),
|
||||
)
|
||||
// 合并大体积上报路由(继承各自的 body limit)
|
||||
.merge(report_router)
|
||||
.layer(DefaultBodyLimit::max(DEFAULT_BODY_LIMIT));
|
||||
|
||||
let api_router = if state.auth_token.is_some() {
|
||||
info!("已为 DCTS 服务端 API 路由启用 Bearer Token / X-API-Key 访问控制鉴权");
|
||||
// 鉴权启用条件:未应急关闭,且配置了 admin 凭据。
|
||||
let auth_enabled = !state.auth_disabled && state.admin_token.is_some();
|
||||
|
||||
let api_router = if auth_enabled {
|
||||
info!("已启用 API 身份鉴权保护(Admin 端点需 admin token 验证;Node 节点免 Token 提交申请,经 Dashboard 管理员审批授权下发)");
|
||||
// 鉴权失败限流(防 token 在线暴力):外层先判 IP 限流,内层再做鉴权。
|
||||
// 限流状态为 20 次/分钟(按 IP),超阈值返回 429。
|
||||
let limiter = api::rate_limit::RateLimiter::new(20, std::time::Duration::from_secs(60));
|
||||
let rate_limit_layer =
|
||||
axum::middleware::from_fn_with_state(limiter, api::rate_limit::rate_limit_middleware);
|
||||
let auth_layer = axum::middleware::from_fn_with_state(state.clone(), api::auth_middleware);
|
||||
api_router.layer(auth_layer)
|
||||
api_router.layer(auth_layer).layer(rate_limit_layer)
|
||||
} else {
|
||||
tracing::warn!("⚠️ 警告:未检测到 DCTS_AUTH_TOKEN 环境变量,服务端目前运行在【内网无鉴权模式】!所有 REST API 接口均为公开可访问状态。");
|
||||
tracing::warn!(
|
||||
"⚠️ 警告:未配置 DCTS_ADMIN_TOKEN / DCTS_ENROLLMENT_TOKEN(且未启用 DCTS_AUTH_DISABLE),\
|
||||
服务端运行在【无鉴权模式】!公网部署务必配置凭据。"
|
||||
);
|
||||
api_router
|
||||
};
|
||||
|
||||
// Host Dashboard SPA static files from dashboard/dist if directory exists or fallback to index.html
|
||||
let serve_dir = ServeDir::new("dashboard/dist")
|
||||
.fallback(ServeFile::new("dashboard/dist/index.html"));
|
||||
let serve_dir =
|
||||
ServeDir::new("dashboard/dist").fallback(ServeFile::new("dashboard/dist/index.html"));
|
||||
|
||||
// 安全响应头(CSP / nosniff / DENY / Referrer-Policy)。
|
||||
let security_headers = axum::middleware::from_fn(security_headers_middleware);
|
||||
|
||||
let app = Router::new()
|
||||
// 独立健康检查端点:不走鉴权、不走 CORS/body 限制,专供 docker healthcheck 与外部监控探测。
|
||||
// 开启鉴权后 /api/status 会返回 401,导致容器被判定不健康而反复重启,故单独提供 /healthz。
|
||||
.route("/healthz", get(api::status::healthz))
|
||||
.nest("/api", api_router)
|
||||
.layer(CorsLayer::permissive())
|
||||
.layer(server::cors::build_cors_layer())
|
||||
.layer(security_headers)
|
||||
.fallback_service(serve_dir)
|
||||
.with_state(state);
|
||||
|
||||
@@ -149,14 +347,53 @@ async fn main() -> Result<()> {
|
||||
info!("DCTS 服务端已在 http://{} 启动监听", addr);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async {
|
||||
let _ = tokio::signal::ctrl_c().await;
|
||||
info!("收到 Ctrl+C 终止信号,DCTS 服务端准备优雅关闭...");
|
||||
})
|
||||
.await?;
|
||||
// into_make_service_with_connect_info:让限流中间件能从连接拿到客户端 IP(反代场景则用 X-Forwarded-For)
|
||||
axum::serve(
|
||||
listener,
|
||||
app.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.with_graceful_shutdown(async {
|
||||
let _ = tokio::signal::ctrl_c().await;
|
||||
info!("收到 Ctrl+C 终止信号,DCTS 服务端准备优雅关闭...");
|
||||
})
|
||||
.await?;
|
||||
|
||||
info!("DCTS 服务端已安全关闭。");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 注入安全响应头的中间件函数。
|
||||
async fn security_headers_middleware(
|
||||
req: axum::http::Request<axum::body::Body>,
|
||||
next: axum::middleware::Next,
|
||||
) -> axum::response::Response {
|
||||
let mut resp = next.run(req).await;
|
||||
|
||||
let headers = resp.headers_mut();
|
||||
// CSP:default-src 'self';放行 Google Fonts(index.html 引用);允许 data: 图片。
|
||||
// 已移除 'unsafe-eval':dashboard 构建产物不使用 eval/new Function(已核实),保留它会
|
||||
// 显著削弱 CSP 的脚本注入防护。'unsafe-inline' 暂留(静态 SPA 内联脚本/handler 需要),
|
||||
// 彻底方案需前端改造为外链 + per-request nonce 注入,见 docs TODO。
|
||||
headers
|
||||
.entry(axum::http::header::CONTENT_SECURITY_POLICY)
|
||||
.or_insert_with(|| {
|
||||
HeaderValue::from_static(
|
||||
"default-src 'self'; script-src 'self' 'unsafe-inline'; \
|
||||
style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; \
|
||||
font-src 'self' data: https://fonts.gstatic.com; \
|
||||
connect-src 'self'; img-src 'self' data: blob:; \
|
||||
frame-ancestors 'none'",
|
||||
)
|
||||
});
|
||||
headers
|
||||
.entry(axum::http::header::X_CONTENT_TYPE_OPTIONS)
|
||||
.or_insert_with(|| HeaderValue::from_static("nosniff"));
|
||||
headers
|
||||
.entry(axum::http::header::X_FRAME_OPTIONS)
|
||||
.or_insert_with(|| HeaderValue::from_static("DENY"));
|
||||
headers
|
||||
.entry(axum::http::HeaderName::from_static("referrer-policy"))
|
||||
.or_insert_with(|| HeaderValue::from_static("strict-origin-when-cross-origin"));
|
||||
|
||||
resp
|
||||
}
|
||||
|
||||
+333
-52
@@ -23,13 +23,32 @@ impl GridScheduler {
|
||||
}
|
||||
}
|
||||
|
||||
/// Expands grid points from config and registers them into the database
|
||||
pub async fn initialize_grid(&self, cfg: &GridConfig) -> Result<()> {
|
||||
if let Err(e) = self.queue.clear_queue().await {
|
||||
tracing::warn!("初始化网格时清理闲置排队记录发生警告: {}", e);
|
||||
/// Expands grid points from config and registers them into the database.
|
||||
///
|
||||
/// 多工作流分区(#3 修复):
|
||||
/// - 仅清理**本工作流**的排队任务(clear_queue_by_workflow),不再 clear_queue() 全局清空,
|
||||
/// 避免启动工作流 B 时误删工作流 A 的在队任务。
|
||||
/// - 仅重置**本工作流**的 queued 点为 pending(reset_queued_grid_points_to_pending 带 wf),
|
||||
/// 避免误伤其他工作流。
|
||||
/// - upsert 带 workflow_name,使同一物理点可属于多个工作流。
|
||||
pub async fn initialize_grid(&self, cfg: &GridConfig, workflow_name: &str) -> Result<()> {
|
||||
if let Err(e) = self.queue.clear_queue_by_workflow(workflow_name).await {
|
||||
tracing::warn!(
|
||||
"初始化工作流 {} 网格时清理该流闲置排队记录发生警告: {}",
|
||||
workflow_name,
|
||||
e
|
||||
);
|
||||
}
|
||||
if let Err(e) = self.db.reset_queued_grid_points_to_pending().await {
|
||||
tracing::warn!("重置网格状态到 pending 处理过程遇到异常: {}", e);
|
||||
if let Err(e) = self
|
||||
.db
|
||||
.reset_queued_grid_points_to_pending(workflow_name)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
"重置工作流 {} 网格状态到 pending 处理过程遇到异常: {}",
|
||||
workflow_name,
|
||||
e
|
||||
);
|
||||
}
|
||||
let mut points = Vec::new();
|
||||
|
||||
@@ -59,9 +78,21 @@ impl GridScheduler {
|
||||
a.cno_sum()
|
||||
.partial_cmp(&b.cno_sum())
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
.then_with(|| a.teff.partial_cmp(&b.teff).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.then_with(|| b.logg.partial_cmp(&a.logg).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.then_with(|| a.loghe.partial_cmp(&b.loghe).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.then_with(|| {
|
||||
a.teff
|
||||
.partial_cmp(&b.teff)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
.then_with(|| {
|
||||
b.logg
|
||||
.partial_cmp(&a.logg)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
.then_with(|| {
|
||||
a.loghe
|
||||
.partial_cmp(&b.loghe)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
});
|
||||
|
||||
// Group into Waves by cno_sum
|
||||
@@ -79,46 +110,100 @@ impl GridScheduler {
|
||||
current_cno = Some(cno);
|
||||
}
|
||||
|
||||
self.db.upsert_grid_point(pt, wave_idx).await?;
|
||||
// upsert 是幂等的 ON CONFLICT DO NOTHING:若 initialize_grid 中途失败,
|
||||
// 重新 start 该工作流会自然补齐(#4 半初始化回退由幂等性消解)。
|
||||
self.db
|
||||
.upsert_grid_point(pt, wave_idx, workflow_name)
|
||||
.await?;
|
||||
}
|
||||
|
||||
info!("已在数据库中成功初始化并记录 {} 个恒星大气网格点", points.len());
|
||||
info!(
|
||||
"已在数据库中成功初始化并记录工作流 {} 的 {} 个恒星大气网格点",
|
||||
workflow_name,
|
||||
points.len()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_active_timeout_sec(&self) -> u64 {
|
||||
if let Ok(yamls) = self.db.get_running_workflow_config_yamls().await {
|
||||
for yaml in yamls {
|
||||
if let Ok(cfg) = serde_yaml::from_str::<GridConfig>(&yaml) {
|
||||
return cfg.timeout_sec;
|
||||
}
|
||||
/// 读取指定工作流的 timeout_sec(按工作流分区:多工作流各有自己的超时配置)。
|
||||
async fn get_workflow_timeout_sec(&self, workflow_name: &str) -> u64 {
|
||||
if let Ok(Some(wf)) = self.db.get_workflow(workflow_name).await {
|
||||
if let Ok(cfg) = serde_yaml::from_str::<GridConfig>(&wf.config_yaml) {
|
||||
return cfg.timeout_sec;
|
||||
}
|
||||
}
|
||||
7200
|
||||
}
|
||||
|
||||
/// Enqueues pending grid points into MQ with active seed detection and batching
|
||||
/// 读取指定工作流的 seed_step_fallback 配置。
|
||||
async fn get_workflow_seed_step_fallback(&self, workflow_name: &str) -> bool {
|
||||
if let Ok(Some(wf)) = self.db.get_workflow(workflow_name).await {
|
||||
if let Ok(cfg) = serde_yaml::from_str::<GridConfig>(&wf.config_yaml) {
|
||||
return cfg.seed_step_fallback;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Enqueues pending grid points into MQ with active seed detection and batching.
|
||||
///
|
||||
/// 多工作流分区(#3 修复):对**每个** running/initializing 工作流分别派发任务,
|
||||
/// 替代原来「全局只一个 running workflow」的 LIMIT 1 假设。各工作流独立 batch、
|
||||
/// 独立 seed 匹配(seeds 仍是全局共享的物理资源池)。
|
||||
pub async fn schedule_pending_tasks(&self) -> Result<usize> {
|
||||
if !self.db.has_running_workflow().await? {
|
||||
let workflows = self.db.get_running_workflow_names().await?;
|
||||
if workflows.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let timeout_sec = self.get_active_timeout_sec().await;
|
||||
let batch_limit: usize = std::env::var("DCTS_BATCH_LIMIT")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(100);
|
||||
|
||||
// SQL 层直接附加 LIMIT = batch_limit 筛选,完全免除数万点位无谓内存反序列化和空耗对象释放开销
|
||||
let pending = self.db.get_pending_grid_points_limit(batch_limit).await?;
|
||||
let mut total_dispatched = 0;
|
||||
|
||||
for wf in &workflows {
|
||||
let dispatched = self
|
||||
.schedule_pending_tasks_for_workflow(wf, batch_limit)
|
||||
.await?;
|
||||
total_dispatched += dispatched;
|
||||
}
|
||||
|
||||
if total_dispatched > 0 {
|
||||
info!(
|
||||
"已成功将 {} 个待计算网格点推进任务队列(跨 {} 个工作流)",
|
||||
total_dispatched,
|
||||
workflows.len()
|
||||
);
|
||||
}
|
||||
|
||||
Ok(total_dispatched)
|
||||
}
|
||||
|
||||
/// 为单个工作流派发 pending 点。
|
||||
async fn schedule_pending_tasks_for_workflow(
|
||||
&self,
|
||||
workflow_name: &str,
|
||||
batch_limit: usize,
|
||||
) -> Result<usize> {
|
||||
let timeout_sec = self.get_workflow_timeout_sec(workflow_name).await;
|
||||
|
||||
// SQL 层直接附加 LIMIT = batch_limit + workflow_name 筛选,完全免除数万点位无谓内存反序列化
|
||||
let pending = self
|
||||
.db
|
||||
.get_pending_grid_points_limit(batch_limit, workflow_name)
|
||||
.await?;
|
||||
let mut dispatched = 0;
|
||||
|
||||
for (name, params, _wave) in pending {
|
||||
|
||||
// Check if any seed is available in DB for active SeedStep scheduling
|
||||
// Check if any seed is available in DB for active SeedStep scheduling(seeds 全局共享)
|
||||
let (task_type, seed_name) = match self.db.find_best_seed_from_db(¶ms).await {
|
||||
Ok(Some(seed_match)) => {
|
||||
info!("网格点 {} 匹配到数据库近邻种子 {} (距离: {:.2}),安排 SeedStep 热启动调度", name, seed_match.name, seed_match.distance);
|
||||
info!(
|
||||
"工作流 {} 网格点 {} 匹配到数据库近邻种子 {} (距离: {:.2}),安排 SeedStep 热启动调度",
|
||||
workflow_name, name, seed_match.name, seed_match.distance
|
||||
);
|
||||
(TaskType::SeedStep, Some(seed_match.name))
|
||||
}
|
||||
_ => (TaskType::ColdRun, None),
|
||||
@@ -131,48 +216,96 @@ impl GridScheduler {
|
||||
task_type,
|
||||
seed_point_name: seed_name,
|
||||
timeout_sec,
|
||||
workflow_name: Some(workflow_name.to_string()),
|
||||
};
|
||||
|
||||
self.db.insert_task(&task_spec).await?;
|
||||
// 采用先标记 DB 状态为 Queued 后发 MQ 的时序,防止推入 MQ 后数据库修改异常导向下一轮误重投
|
||||
self.db.update_grid_status(&name, common::models::GridPointStatus::Queued).await?;
|
||||
self.db
|
||||
.update_grid_status(
|
||||
&name,
|
||||
common::models::GridPointStatus::Queued,
|
||||
workflow_name,
|
||||
)
|
||||
.await?;
|
||||
match self.queue.push_task(&task_spec).await {
|
||||
Ok(_) => {
|
||||
dispatched += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("将任务 {} 推入 MQ 队列失败,回滚网格点状态: {}", name, e);
|
||||
let _ = self.db.update_grid_status(&name, common::models::GridPointStatus::Pending).await;
|
||||
tracing::warn!(
|
||||
"将任务 {} 推入 MQ 队列失败,执行严格状态回滚以避免脏数据: {}",
|
||||
name,
|
||||
e
|
||||
);
|
||||
if let Err(db_e) = self
|
||||
.db
|
||||
.update_grid_status(
|
||||
&name,
|
||||
common::models::GridPointStatus::Pending,
|
||||
workflow_name,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::error!(
|
||||
"关键性回滚异常:任务 {} 无法重置回 Pending: {}",
|
||||
name,
|
||||
db_e
|
||||
);
|
||||
}
|
||||
let _ = self.queue.remove_task(&task_spec.task_id.to_string()).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if dispatched > 0 {
|
||||
info!("已成功将 {} 个待计算网格点推进任务队列", dispatched);
|
||||
}
|
||||
|
||||
Ok(dispatched)
|
||||
}
|
||||
|
||||
/// Triggers seed_step fallback for a failed point if a seed is available
|
||||
pub async fn trigger_seed_step_fallback(&self, params: &GridPointParams) -> Result<bool> {
|
||||
if !self.db.has_running_workflow().await? {
|
||||
/// Triggers seed_step fallback for a failed point if a seed is available.
|
||||
///
|
||||
/// 多工作流分区(#3 修复):传入 `workflow_name` 明确该失败点所属工作流,
|
||||
/// 用该工作流自身的 timeout / seed_step_fallback 配置,并把 TaskSpec.workflow_name
|
||||
/// 绑定到该工作流。
|
||||
///
|
||||
/// 语义(种子回退仅一次):
|
||||
/// - 仅当该工作流配置 `seed_step_fallback: true` 时才考虑回退;
|
||||
/// - 仅当该点**尚未**派发过任何 seed_step 任务时才回退一次;
|
||||
/// - 找不到合适近邻种子则不回退,由调用方保持 failed 终态。
|
||||
pub async fn trigger_seed_step_fallback(
|
||||
&self,
|
||||
params: &GridPointParams,
|
||||
workflow_name: &str,
|
||||
) -> Result<bool> {
|
||||
// 该工作流须仍处于 running 态才回退(避免 stop 后继续派发)
|
||||
let still_running = self
|
||||
.db
|
||||
.get_running_workflow_names()
|
||||
.await?
|
||||
.iter()
|
||||
.any(|w| w == workflow_name);
|
||||
if !still_running {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
if !self.get_workflow_seed_step_fallback(workflow_name).await {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let name = params.model_name();
|
||||
if let Ok(Some((status, attempt_count))) = self.db.get_grid_point_status(&name).await {
|
||||
if status == "failed" || attempt_count >= 3 {
|
||||
info!("网格点 {} 已达到最大重试次数 ({}) 或处于 failed 状态,跳过种子热启动回退", name, attempt_count);
|
||||
return Ok(false);
|
||||
}
|
||||
// 种子回退仅一次:该点在该工作流中已经派发过 seed_step 任务就不再触发新的回退
|
||||
if self.db.has_seed_step_attempt(&name, workflow_name).await? {
|
||||
info!(
|
||||
"网格点 {} 已使用过一次种子热启动回退,不再重复回退,保持 failed 终态",
|
||||
name
|
||||
);
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
// seeds 全局共享:跨工作流复用已收敛的邻近种子
|
||||
let seed_match_opt = self.db.find_best_seed_from_db(params).await.ok().flatten();
|
||||
|
||||
if let Some(seed_match) = seed_match_opt {
|
||||
let timeout_sec = self.get_active_timeout_sec().await;
|
||||
let timeout_sec = self.get_workflow_timeout_sec(workflow_name).await;
|
||||
let name = params.model_name();
|
||||
let task_spec = TaskSpec {
|
||||
task_id: Uuid::new_v4(),
|
||||
@@ -181,22 +314,38 @@ impl GridScheduler {
|
||||
task_type: TaskType::SeedStep,
|
||||
seed_point_name: Some(seed_match.name.clone()),
|
||||
timeout_sec,
|
||||
workflow_name: Some(workflow_name.to_string()),
|
||||
};
|
||||
|
||||
self.db.insert_task(&task_spec).await?;
|
||||
self.db.update_grid_status(&name, common::models::GridPointStatus::Queued).await?;
|
||||
self.db
|
||||
.update_grid_status(
|
||||
&name,
|
||||
common::models::GridPointStatus::Queued,
|
||||
workflow_name,
|
||||
)
|
||||
.await?;
|
||||
if let Err(e) = self.queue.push_task(&task_spec).await {
|
||||
let _ = self.db.update_grid_status(&name, common::models::GridPointStatus::Pending).await;
|
||||
let _ = self
|
||||
.db
|
||||
.update_grid_status(
|
||||
&name,
|
||||
common::models::GridPointStatus::Pending,
|
||||
workflow_name,
|
||||
)
|
||||
.await;
|
||||
let _ = self.queue.remove_task(&task_spec.task_id.to_string()).await;
|
||||
return Err(e);
|
||||
}
|
||||
info!("触发种子步进 (seed_step):网格点 {} 将使用 6 维近邻种子 {} 热启动重试", name, seed_match.name);
|
||||
info!(
|
||||
"触发种子步进 (seed_step):工作流 {} 网格点 {} 将使用 6 维近邻种子 {} 热启动重试",
|
||||
workflow_name, name, seed_match.name
|
||||
);
|
||||
Ok(true)
|
||||
} else {
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -212,8 +361,16 @@ mod tests {
|
||||
let results_dir = temp_dir.path().join("results");
|
||||
|
||||
let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
|
||||
let queue = Arc::new(SqliteTaskQueue::new(&queue_db_path.to_string_lossy()).await.unwrap());
|
||||
let scheduler = GridScheduler::new(db.clone(), queue.clone(), results_dir.to_string_lossy().to_string());
|
||||
let queue = Arc::new(
|
||||
SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let scheduler = GridScheduler::new(
|
||||
db.clone(),
|
||||
queue.clone(),
|
||||
results_dir.to_string_lossy().to_string(),
|
||||
);
|
||||
|
||||
let cfg = GridConfig {
|
||||
grid: GridAxesConfig {
|
||||
@@ -238,17 +395,141 @@ mod tests {
|
||||
linelist: None,
|
||||
};
|
||||
|
||||
scheduler.initialize_grid(&cfg).await.unwrap();
|
||||
db.upsert_workflow("test_wf", None, "", "running").await.unwrap();
|
||||
scheduler.initialize_grid(&cfg, "test_wf").await.unwrap();
|
||||
db.upsert_workflow("test_wf", None, "", "running")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let pending = db.get_pending_grid_points().await.unwrap();
|
||||
let pending = db.get_pending_grid_points("test_wf").await.unwrap();
|
||||
assert_eq!(pending.len(), 1);
|
||||
|
||||
let dispatched = scheduler.schedule_pending_tasks().await.unwrap();
|
||||
assert_eq!(dispatched, 1);
|
||||
|
||||
let popped = queue.pop_task().await.unwrap();
|
||||
let popped = queue.pop_task("test-node").await.unwrap();
|
||||
assert!(popped.is_some());
|
||||
}
|
||||
}
|
||||
|
||||
/// 多工作流分区调度测试(#3 修复验证):
|
||||
/// 1. wf_a 调度推入队列的任务,在初始化 wf_b 后依然存在(initialize_grid 改用
|
||||
/// clear_queue_by_workflow,不再全局 clear_queue)。
|
||||
/// 2. 两个 running 工作流的 pending 点都能被 schedule_pending_tasks 派发。
|
||||
#[tokio::test]
|
||||
async fn test_multi_workflow_dispatch_isolation() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let db = Database::new(&temp_dir.path().join("mw_db.db").to_string_lossy())
|
||||
.await
|
||||
.unwrap();
|
||||
let queue = Arc::new(
|
||||
SqliteTaskQueue::new(&temp_dir.path().join("mw_queue.db").to_string_lossy())
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let scheduler = GridScheduler::new(db.clone(), queue.clone(), "results".to_string());
|
||||
|
||||
let mk_cfg = |teff: f64| GridConfig {
|
||||
grid: GridAxesConfig {
|
||||
teff: vec![teff],
|
||||
logg: vec![5.5],
|
||||
loghe: vec![-1.0],
|
||||
logc: vec![-2.0],
|
||||
logn: vec![-2.0],
|
||||
logo: vec![-2.0],
|
||||
},
|
||||
chain: vec![],
|
||||
synspec: None,
|
||||
nworkers: 4,
|
||||
timeout_sec: 3600,
|
||||
resume: true,
|
||||
seed_step_fallback: true,
|
||||
results: None,
|
||||
itek_fallback: vec![],
|
||||
niter: Some(100),
|
||||
template: None,
|
||||
fort55: None,
|
||||
linelist: None,
|
||||
};
|
||||
|
||||
// wf_a 初始化并推入队列
|
||||
scheduler
|
||||
.initialize_grid(&mk_cfg(35000.0), "wf_a")
|
||||
.await
|
||||
.unwrap();
|
||||
db.upsert_workflow("wf_a", None, "", "running")
|
||||
.await
|
||||
.unwrap();
|
||||
let d_a = scheduler.schedule_pending_tasks().await.unwrap();
|
||||
assert_eq!(d_a, 1);
|
||||
// 任务已在队
|
||||
assert!(queue.pop_task("node-a").await.unwrap().is_some());
|
||||
|
||||
// 重新推一个 wf_a 任务(上一行 pop 掉了),再初始化 wf_b
|
||||
db.update_grid_status(
|
||||
&GridPointParams {
|
||||
teff: 35000.0,
|
||||
logg: 5.5,
|
||||
loghe: -1.0,
|
||||
logc: -2.0,
|
||||
logn: -2.0,
|
||||
logo: -2.0,
|
||||
}
|
||||
.model_name(),
|
||||
common::models::GridPointStatus::Pending,
|
||||
"wf_a",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _ = scheduler
|
||||
.schedule_pending_tasks_for_workflow("wf_a", 100)
|
||||
.await
|
||||
.unwrap();
|
||||
// 此时 wf_a 队列里应有一个任务
|
||||
assert_eq!(
|
||||
queue
|
||||
.pop_task("node-a")
|
||||
.await
|
||||
.unwrap()
|
||||
.and_then(|t| t.workflow_name),
|
||||
Some("wf_a".to_string())
|
||||
);
|
||||
|
||||
// 关键断言:把 wf_a 任务重新推回队列后,初始化 wf_b 不应清空它。
|
||||
db.update_grid_status(
|
||||
&GridPointParams {
|
||||
teff: 35000.0,
|
||||
logg: 5.5,
|
||||
loghe: -1.0,
|
||||
logc: -2.0,
|
||||
logn: -2.0,
|
||||
logo: -2.0,
|
||||
}
|
||||
.model_name(),
|
||||
common::models::GridPointStatus::Pending,
|
||||
"wf_a",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _ = scheduler
|
||||
.schedule_pending_tasks_for_workflow("wf_a", 100)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// 初始化 wf_b(内部 clear_queue_by_workflow("wf_b"),不该动 wf_a 的任务)
|
||||
scheduler
|
||||
.initialize_grid(&mk_cfg(40000.0), "wf_b")
|
||||
.await
|
||||
.unwrap();
|
||||
db.upsert_workflow("wf_b", None, "", "running")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// wf_a 的任务仍在队:可被 node 弹出,且 workflow_name == wf_a
|
||||
let popped_a = queue.pop_task("node-a").await.unwrap();
|
||||
assert!(popped_a.is_some(), "初始化 wf_b 不应清空 wf_a 的队列任务");
|
||||
assert_eq!(popped_a.unwrap().workflow_name, Some("wf_a".to_string()));
|
||||
|
||||
// wf_b 的点也能被调度(两个 running 工作流并存)
|
||||
let d_b = scheduler.schedule_pending_tasks().await.unwrap();
|
||||
assert!(d_b >= 1, "wf_b 的 pending 点应被派发");
|
||||
}
|
||||
}
|
||||
|
||||
+1175
-17
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user