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:
@@ -1,2 +1 @@
|
||||
pub mod sqlite_queue;
|
||||
|
||||
|
||||
+219
-27
@@ -10,11 +10,32 @@ struct SqliteCustomizer;
|
||||
|
||||
impl r2d2::CustomizeConnection<rusqlite::Connection, rusqlite::Error> for SqliteCustomizer {
|
||||
fn on_acquire(&self, conn: &mut rusqlite::Connection) -> Result<(), rusqlite::Error> {
|
||||
conn.pragma_update(None, "busy_timeout", 5000)?;
|
||||
// 与主库一致:高并发 claim/report 下给 SQLITE_BUSY 足够重试窗口。
|
||||
conn.pragma_update(None, "busy_timeout", 15000)?;
|
||||
conn.pragma_update(None, "wal_autocheckpoint", 1000)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 SQLite db 文件及其 WAL/SHM 侧车文件权限收紧为 0600(仅 owner 读写)。
|
||||
/// 与主库 dcts.db 的口径一致,作为纵深防御(队列库不含 token,但含任务 payload)。
|
||||
#[cfg(unix)]
|
||||
fn restrict_db_file_perms(db_path: &str) {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let candidates = [
|
||||
std::path::PathBuf::from(db_path),
|
||||
std::path::PathBuf::from(format!("{}-wal", db_path)),
|
||||
std::path::PathBuf::from(format!("{}-shm", db_path)),
|
||||
];
|
||||
for p in candidates {
|
||||
if let Ok(meta) = std::fs::metadata(&p) {
|
||||
let mut perms = meta.permissions();
|
||||
perms.set_mode(0o600);
|
||||
let _ = std::fs::set_permissions(&p, perms);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct SqliteTaskQueue {
|
||||
pool: Pool<SqliteConnectionManager>,
|
||||
@@ -29,11 +50,15 @@ impl SqliteTaskQueue {
|
||||
}
|
||||
let manager = SqliteConnectionManager::file(&db_path_owned);
|
||||
let pool = Pool::builder()
|
||||
.max_size(4)
|
||||
.max_size(8)
|
||||
.connection_customizer(Box::new(SqliteCustomizer))
|
||||
.build(manager)
|
||||
.context("Failed to build SQLite queue connection pool")?;
|
||||
|
||||
// 收紧队列 db 文件权限为 0600(与主库口径一致,纵深防御)
|
||||
#[cfg(unix)]
|
||||
restrict_db_file_perms(&db_path_owned);
|
||||
|
||||
let conn = pool.get()?;
|
||||
let _: String = conn.pragma_update_and_check(None, "journal_mode", "WAL", |r| r.get(0))?;
|
||||
conn.execute(
|
||||
@@ -42,14 +67,39 @@ impl SqliteTaskQueue {
|
||||
payload TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL,
|
||||
claimed_at DATETIME
|
||||
claimed_at DATETIME,
|
||||
workflow_name TEXT,
|
||||
claimed_by_node_id TEXT
|
||||
)",
|
||||
[],
|
||||
)?;
|
||||
// 兼容旧库:若 task_queue 表已存在但缺少 workflow_name / claimed_by_node_id 列,则补列。
|
||||
// SQLite 的 ALTER TABLE ADD COLUMN 是在线操作,旧数据该列默认 NULL。
|
||||
// PRAGMA table_info 检测列是否存在以实现幂等 migration。
|
||||
let has_col = |conn: &rusqlite::Connection, col: &str| -> rusqlite::Result<bool> {
|
||||
let mut stmt = conn.prepare("PRAGMA table_info(task_queue)")?;
|
||||
let rows = stmt.query_map([], |r| r.get::<_, String>(1))?;
|
||||
for r in rows {
|
||||
if r.map(|name| name == col).unwrap_or(false) {
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
Ok(false)
|
||||
};
|
||||
if !has_col(&conn, "workflow_name")? {
|
||||
conn.execute("ALTER TABLE task_queue ADD COLUMN workflow_name TEXT", [])?;
|
||||
}
|
||||
if !has_col(&conn, "claimed_by_node_id")? {
|
||||
conn.execute("ALTER TABLE task_queue ADD COLUMN claimed_by_node_id TEXT", [])?;
|
||||
}
|
||||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_task_queue_status_created ON task_queue(status, created_at)",
|
||||
[],
|
||||
)?;
|
||||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_task_queue_workflow ON task_queue(workflow_name)",
|
||||
[],
|
||||
)?;
|
||||
Ok(pool)
|
||||
})
|
||||
.await??;
|
||||
@@ -61,14 +111,15 @@ impl SqliteTaskQueue {
|
||||
pub async fn push_task(&self, task: &TaskSpec) -> Result<()> {
|
||||
let payload = serde_json::to_string(task)?;
|
||||
let task_id_str = task.task_id.to_string();
|
||||
let workflow_name = task.workflow_name.clone();
|
||||
let pool = self.pool.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO task_queue (task_id, payload, status, created_at)
|
||||
VALUES (?1, ?2, 'pending', datetime('now'))",
|
||||
params![task_id_str, payload],
|
||||
"INSERT OR REPLACE INTO task_queue (task_id, payload, status, created_at, workflow_name)
|
||||
VALUES (?1, ?2, 'pending', datetime('now'), ?3)",
|
||||
params![task_id_str, payload, workflow_name],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
@@ -77,8 +128,9 @@ impl SqliteTaskQueue {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn pop_task(&self) -> Result<Option<TaskSpec>> {
|
||||
pub async fn pop_task(&self, claimant_node_id: &str) -> Result<Option<TaskSpec>> {
|
||||
let pool = self.pool.clone();
|
||||
let claimant = claimant_node_id.to_string();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<Option<TaskSpec>> {
|
||||
let mut attempts = 0;
|
||||
@@ -109,9 +161,11 @@ impl SqliteTaskQueue {
|
||||
|
||||
let task: TaskSpec = serde_json::from_str(&payload)?;
|
||||
|
||||
// 记录任务归属:claim 时写入领用方 node_id,供 report 阶段校验,
|
||||
// 杜绝「节点 A 领用、节点 B 上报」的跨节点伪造结果投毒。
|
||||
tx.execute(
|
||||
"UPDATE task_queue SET status = 'claimed', claimed_at = datetime('now') WHERE task_id = ?1",
|
||||
params![task_id],
|
||||
"UPDATE task_queue SET status = 'claimed', claimed_at = datetime('now'), claimed_by_node_id = ?2 WHERE task_id = ?1",
|
||||
params![task_id, claimant],
|
||||
)?;
|
||||
|
||||
tx.commit()?;
|
||||
@@ -139,8 +193,13 @@ impl SqliteTaskQueue {
|
||||
let id_owned = task_id.to_string();
|
||||
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
conn.execute("DELETE FROM task_queue WHERE task_id = ?1", params![id_owned])?;
|
||||
let conn = pool
|
||||
.get()
|
||||
.map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
conn.execute(
|
||||
"DELETE FROM task_queue WHERE task_id = ?1",
|
||||
params![id_owned],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
@@ -148,13 +207,60 @@ impl SqliteTaskQueue {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn requeue_stale_tasks(&self, stale_sec: u64) -> Result<Vec<String>> {
|
||||
/// 校验指定 task 是否由指定 node 领用(claimed 态且 claimed_by_node_id 匹配)。
|
||||
///
|
||||
/// 用于 report_task 阶段防止跨节点伪造结果:只有真正领用该 task 的 node 才能上报结果。
|
||||
/// 返回 (point_name, workflow_name):匹配时附带二者供 report 进一步校验「上报的点与领用的
|
||||
/// task 一致」并把 workflow_name 传给 record_task_report 以定向更新对应工作流的 grid_points
|
||||
/// (多工作流分区:避免按 name 全局更新误改其他工作流同名点)。
|
||||
pub async fn verify_task_claim(
|
||||
&self,
|
||||
task_id: &str,
|
||||
claimant_node_id: &str,
|
||||
) -> Result<Option<(String, Option<String>)>> {
|
||||
let pool = self.pool.clone();
|
||||
let task_id = task_id.to_string();
|
||||
let claimant = claimant_node_id.to_string();
|
||||
|
||||
let res =
|
||||
tokio::task::spawn_blocking(move || -> Result<Option<(String, Option<String>)>> {
|
||||
let conn = pool
|
||||
.get()
|
||||
.map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
// 仅 claimed 态(尚未被 report 清理)且归属匹配才算有效领用
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT payload FROM task_queue
|
||||
WHERE task_id = ?1 AND claimed_by_node_id = ?2 AND status = 'claimed' LIMIT 1",
|
||||
)?;
|
||||
let row = stmt.query_row(params![task_id, claimant], |r| r.get::<_, String>(0));
|
||||
match row {
|
||||
Ok(payload) => {
|
||||
// 解析 payload 取出 point_name + workflow_name,供调用方校验与定向更新
|
||||
let task: TaskSpec = serde_json::from_str(&payload)?;
|
||||
Ok(Some((task.point_name, task.workflow_name)))
|
||||
}
|
||||
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
})
|
||||
.await??;
|
||||
Ok(res)
|
||||
}
|
||||
|
||||
/// 重投超时 claimed 任务回 pending,返回每个被重投任务的 (point_name, workflow_name)。
|
||||
///
|
||||
/// 返回 workflow_name 供调用方(main.rs 后台循环)按工作流分组调用
|
||||
/// reset_specific_grid_points_to_pending,避免跨工作流误改同名点(多工作流分区)。
|
||||
pub async fn requeue_stale_tasks(
|
||||
&self,
|
||||
stale_sec: u64,
|
||||
) -> Result<Vec<(String, Option<String>)>> {
|
||||
let pool = self.pool.clone();
|
||||
|
||||
let names = tokio::task::spawn_blocking(move || -> Result<Vec<String>> {
|
||||
let entries = tokio::task::spawn_blocking(move || -> Result<Vec<(String, Option<String>)>> {
|
||||
let mut conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
let tx = conn.transaction()?;
|
||||
let mut point_names = Vec::new();
|
||||
let mut entries = Vec::new();
|
||||
|
||||
{
|
||||
// 改写为单一原子更新带 RETURNING 返回语句,消弭 TOCTOU (Time-Of-Check-To-Time-Of-Use) 竞态问题
|
||||
@@ -165,32 +271,56 @@ impl SqliteTaskQueue {
|
||||
)?;
|
||||
let rows = stmt.query_map(params![stale_sec as i64], |row| row.get::<_, String>(0))?;
|
||||
for r in rows {
|
||||
if let Ok(payload) = r {
|
||||
if let Ok(task) = serde_json::from_str::<TaskSpec>(&payload) {
|
||||
point_names.push(task.point_name);
|
||||
}
|
||||
let payload = match r {
|
||||
Ok(p) => p,
|
||||
Err(_) => continue,
|
||||
};
|
||||
if let Ok(task) = serde_json::from_str::<TaskSpec>(&payload) {
|
||||
entries.push((task.point_name, task.workflow_name));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tx.commit()?;
|
||||
Ok(point_names)
|
||||
Ok(entries)
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(names)
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
pub async fn clear_queue(&self) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool.get().map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
let conn = pool
|
||||
.get()
|
||||
.map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
conn.execute("DELETE FROM task_queue", [])?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 仅清理指定工作流的排队任务。
|
||||
///
|
||||
/// 用于 stop_workflow 按工作流隔离清理,避免在多工作流场景下误清其他工作流的任务。
|
||||
pub async fn clear_queue_by_workflow(&self, workflow_name: &str) -> Result<()> {
|
||||
let pool = self.pool.clone();
|
||||
let wf_owned = workflow_name.to_string();
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
let conn = pool
|
||||
.get()
|
||||
.map_err(|e| anyhow::anyhow!("Queue DB pool error: {}", e))?;
|
||||
conn.execute(
|
||||
"DELETE FROM task_queue WHERE workflow_name = ?1",
|
||||
params![wf_owned],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -203,9 +333,11 @@ mod tests {
|
||||
async fn test_sqlite_task_queue_operations() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let db_path = temp_dir.path().join("test_queue.db");
|
||||
let queue = SqliteTaskQueue::new(&db_path.to_string_lossy()).await.unwrap();
|
||||
let queue = SqliteTaskQueue::new(&db_path.to_string_lossy())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(queue.pop_task().await.unwrap().is_none());
|
||||
assert!(queue.pop_task("test-node").await.unwrap().is_none());
|
||||
|
||||
let task_id = Uuid::new_v4();
|
||||
let task = TaskSpec {
|
||||
@@ -222,24 +354,84 @@ mod tests {
|
||||
task_type: TaskType::ColdRun,
|
||||
seed_point_name: None,
|
||||
timeout_sec: 3600,
|
||||
workflow_name: Some("test_wf".to_string()),
|
||||
};
|
||||
queue.push_task(&task).await.unwrap();
|
||||
|
||||
let popped = queue.pop_task().await.unwrap();
|
||||
let popped = queue.pop_task("test-node").await.unwrap();
|
||||
assert!(popped.is_some());
|
||||
let popped_task = popped.unwrap();
|
||||
assert_eq!(popped_task.task_id, task_id);
|
||||
assert_eq!(popped_task.point_name, task.point_name);
|
||||
|
||||
assert!(queue.pop_task().await.unwrap().is_none());
|
||||
assert!(queue.pop_task("test-node").await.unwrap().is_none());
|
||||
|
||||
let requeued = queue.requeue_stale_tasks(0).await.unwrap();
|
||||
assert_eq!(requeued.len(), 1);
|
||||
|
||||
let popped2 = queue.pop_task().await.unwrap();
|
||||
let popped2 = queue.pop_task("test-node").await.unwrap();
|
||||
assert!(popped2.is_some());
|
||||
|
||||
queue.remove_task(&task_id.to_string()).await.unwrap();
|
||||
assert!(queue.pop_task().await.unwrap().is_none());
|
||||
assert!(queue.pop_task("test-node").await.unwrap().is_none());
|
||||
}
|
||||
|
||||
/// 任务归属校验:领用方 node 匹配才放行,其他 node 校验失败(防跨节点伪造结果)。
|
||||
#[tokio::test]
|
||||
async fn test_verify_task_claim_ownership() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let db_path = temp_dir.path().join("claim_test.db");
|
||||
let queue = SqliteTaskQueue::new(&db_path.to_string_lossy())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let task_id = Uuid::new_v4();
|
||||
let params = GridPointParams {
|
||||
teff: 35000.0,
|
||||
logg: 5.5,
|
||||
loghe: -1.0,
|
||||
logc: -2.0,
|
||||
logn: -2.0,
|
||||
logo: -2.0,
|
||||
};
|
||||
let task = TaskSpec {
|
||||
task_id,
|
||||
point_name: params.model_name(),
|
||||
params: params.clone(),
|
||||
task_type: TaskType::ColdRun,
|
||||
seed_point_name: None,
|
||||
timeout_sec: 60,
|
||||
workflow_name: None,
|
||||
};
|
||||
queue.push_task(&task).await.unwrap();
|
||||
|
||||
// node-A 领用
|
||||
let popped = queue.pop_task("node-A").await.unwrap();
|
||||
assert!(popped.is_some());
|
||||
|
||||
// node-A 校验:匹配,返回绑定的 (point_name, workflow_name)
|
||||
let claim = queue
|
||||
.verify_task_claim(&task_id.to_string(), "node-A")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
claim.map(|(p, _)| p).as_deref(),
|
||||
Some(params.model_name().as_str())
|
||||
);
|
||||
|
||||
// node-B 校验:非领用方,返回 None
|
||||
let claim_b = queue
|
||||
.verify_task_claim(&task_id.to_string(), "node-B")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(claim_b.is_none());
|
||||
|
||||
// 任务被清理(remove)后,任何 node 校验都失败
|
||||
queue.remove_task(&task_id.to_string()).await.unwrap();
|
||||
let claim_after = queue
|
||||
.verify_task_claim(&task_id.to_string(), "node-A")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(claim_after.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user