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:
fmq
2026-07-28 21:54:02 +08:00
parent 4b4238d702
commit b91f1e4fa5
59 changed files with 7697 additions and 1490 deletions
+112 -28
View File
@@ -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,
}
// 手写 Debugtoken 类字段脱敏为 ***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);
}
}
}
+22 -13
View File
@@ -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));
}
}
+26 -7
View File
@@ -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
);
}
}
}
+8 -2
View File
@@ -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
View File
@@ -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 {
-1
View File
@@ -8,4 +8,3 @@ pub mod models;
pub mod nst_writer;
pub mod runner;
pub mod seed_finder;
+5 -3
View File
@@ -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 {
+9 -4
View File
@@ -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);
}
}
-1
View File
@@ -63,4 +63,3 @@ mod tests {
assert!(content.contains("IELCOR=-1"));
}
}
+72 -16
View File
@@ -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);
}
}
+5 -4
View File
@@ -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
View File
@@ -1,2 +1 @@
pub mod sqlite_queue;
+219 -27
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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/403token 失效或被吊销。与 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);
}
+3
View File
@@ -25,5 +25,8 @@ chrono.workspace = true
uuid.workspace = true
tempfile.workspace = true
dotenvy.workspace = true
sha2.workspace = true
hex.workspace = true
subtle = "2"
+154
View File
@@ -0,0 +1,154 @@
//! 管理 APIAdmin 角色)。
//!
//! 提供 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()),
}
}
+121
View File
@@ -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_middlewareRole::Admin)之后,只要到达此 handler 说明 Token 校验必定成功。
pub async fn check_auth() -> impl IntoResponse {
(
StatusCode::OK,
Json(serde_json::json!({
"success": true,
"message": "Token 验证有效",
"authenticated": true
})),
)
}
+73 -20
View File
@@ -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(),
}
}
+51
View File
@@ -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
View File
@@ -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)。
///
/// 下游 handlerheartbeat / 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_idtoken 只存 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_idDB 只存 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()
}
+135 -8
View File
@@ -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()),
}
}
+196
View File
@@ -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
}
+22 -20
View File
@@ -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())
}
+22 -6
View File
@@ -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
View File
@@ -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(&params, &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(&params, &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(&params).await {
if let Err(e) = state
.scheduler
.trigger_seed_step_fallback(&params, &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
View File
@@ -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()),
}
}
+58
View File
@@ -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 策略:仅允许同源或本地 Originlocalhost / 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
View File
File diff suppressed because it is too large Load Diff
+1 -1
View File
@@ -1,4 +1,4 @@
pub mod api;
pub mod cors;
pub mod db;
pub mod scheduler;
+280 -43
View File
@@ -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_dirDCTS_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();
// CSPdefault-src 'self';放行 Google Fontsindex.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
View File
@@ -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 点为 pendingreset_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 schedulingseeds 全局共享)
let (task_type, seed_name) = match self.db.find_best_seed_from_db(&params).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 点应被派发");
}
}
File diff suppressed because it is too large Load Diff