forked from Eeveid/lightOps
实现 LightOps 运维面板基础功能
This commit is contained in:
226
crates/lightops-server/src/task.rs
Normal file
226
crates/lightops-server/src/task.rs
Normal file
@@ -0,0 +1,226 @@
|
||||
use crate::{
|
||||
error::AppError,
|
||||
state::{AppState, TaskReply},
|
||||
};
|
||||
use chrono::Utc;
|
||||
use lightops_common::protocol::ServerMessage;
|
||||
use serde_json::Value;
|
||||
use tokio::sync::oneshot;
|
||||
use uuid::Uuid;
|
||||
|
||||
pub async fn request_agent_task(
|
||||
state: &AppState,
|
||||
agent_id: &str,
|
||||
user_id: Option<i64>,
|
||||
action: &str,
|
||||
params: Value,
|
||||
) -> Result<TaskReply, AppError> {
|
||||
request_agent_task_with_timeout(
|
||||
state,
|
||||
agent_id,
|
||||
user_id,
|
||||
action,
|
||||
params,
|
||||
state.cfg.task_timeout_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn request_agent_task_with_timeout(
|
||||
state: &AppState,
|
||||
agent_id: &str,
|
||||
user_id: Option<i64>,
|
||||
action: &str,
|
||||
params: Value,
|
||||
timeout_secs: u64,
|
||||
) -> Result<TaskReply, AppError> {
|
||||
request_agent_task_with_timeout_and_stored_params(
|
||||
state,
|
||||
agent_id,
|
||||
user_id,
|
||||
action,
|
||||
params.clone(),
|
||||
params,
|
||||
timeout_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn request_agent_task_with_timeout_and_stored_params(
|
||||
state: &AppState,
|
||||
agent_id: &str,
|
||||
user_id: Option<i64>,
|
||||
action: &str,
|
||||
params: Value,
|
||||
stored_params: Value,
|
||||
timeout_secs: u64,
|
||||
) -> Result<TaskReply, AppError> {
|
||||
let task_id = Uuid::new_v4().to_string();
|
||||
sqlx::query(
|
||||
"INSERT INTO tasks (id, agent_id, user_id, action, params_json, status, started_at) VALUES (?, ?, ?, ?, ?, 'running', ?)",
|
||||
)
|
||||
.bind(&task_id)
|
||||
.bind(agent_id)
|
||||
.bind(user_id)
|
||||
.bind(action)
|
||||
.bind(stored_params.to_string())
|
||||
.bind(Utc::now().to_rfc3339())
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
append_task_event(
|
||||
state,
|
||||
&task_id,
|
||||
"info",
|
||||
"任务已创建",
|
||||
Some(&serde_json::json!({ "action": action })),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let handle = {
|
||||
let agents = state.agents.read().await;
|
||||
agents.get(agent_id).cloned()
|
||||
}
|
||||
.ok_or(AppError::AgentOffline)?;
|
||||
|
||||
let (tx, rx) = oneshot::channel();
|
||||
state.pending.insert(task_id.clone(), tx);
|
||||
state
|
||||
.pending_agents
|
||||
.insert(task_id.clone(), agent_id.to_string());
|
||||
let msg = ServerMessage::TaskRequest {
|
||||
task_id: task_id.clone(),
|
||||
action: action.to_string(),
|
||||
params,
|
||||
};
|
||||
if handle.tx.send(msg).is_err() {
|
||||
state.pending.remove(&task_id);
|
||||
state.pending_agents.remove(&task_id);
|
||||
mark_task(state, &task_id, "failed", None, Some("Agent 已断开连接")).await?;
|
||||
return Err(AppError::AgentOffline);
|
||||
}
|
||||
append_task_event(state, &task_id, "info", "任务已下发到 Agent", None).await?;
|
||||
|
||||
let reply = match tokio::time::timeout(std::time::Duration::from_secs(timeout_secs), rx).await {
|
||||
Ok(Ok(reply)) => reply,
|
||||
Ok(Err(_)) => {
|
||||
state.pending_agents.remove(&task_id);
|
||||
mark_task(state, &task_id, "failed", None, Some("Agent 已断开连接")).await?;
|
||||
return Err(AppError::AgentOffline);
|
||||
}
|
||||
Err(_) => {
|
||||
state.pending.remove(&task_id);
|
||||
state.pending_agents.remove(&task_id);
|
||||
mark_task(state, &task_id, "timeout", None, Some("任务超时")).await?;
|
||||
return Err(AppError::Timeout);
|
||||
}
|
||||
};
|
||||
|
||||
let status = if reply.success { "success" } else { "failed" };
|
||||
mark_task(
|
||||
state,
|
||||
&task_id,
|
||||
status,
|
||||
Some(&reply.data),
|
||||
reply.error.as_deref(),
|
||||
)
|
||||
.await?;
|
||||
Ok(reply)
|
||||
}
|
||||
|
||||
pub async fn append_task_event(
|
||||
state: &AppState,
|
||||
task_id: &str,
|
||||
level: &str,
|
||||
message: &str,
|
||||
data: Option<&Value>,
|
||||
) -> Result<(), AppError> {
|
||||
sqlx::query(
|
||||
"INSERT INTO task_events(task_id, level, message, data_json, created_at) VALUES(?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(level)
|
||||
.bind(message.chars().take(500).collect::<String>())
|
||||
.bind(data.map(|value| value.to_string()))
|
||||
.bind(Utc::now().to_rfc3339())
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn fail_agent_pending_tasks(state: &AppState, agent_id: &str, reason: &str) {
|
||||
let task_ids: Vec<String> = state
|
||||
.pending_agents
|
||||
.iter()
|
||||
.filter(|item| item.value() == agent_id)
|
||||
.map(|item| item.key().clone())
|
||||
.collect();
|
||||
|
||||
for task_id in task_ids {
|
||||
state.pending_agents.remove(&task_id);
|
||||
if let Some((_, tx)) = state.pending.remove(&task_id) {
|
||||
let _ = tx.send(TaskReply {
|
||||
success: false,
|
||||
data: serde_json::json!({}),
|
||||
error: Some(reason.to_string()),
|
||||
});
|
||||
}
|
||||
let _ = mark_task(state, &task_id, "failed", None, Some(reason)).await;
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn mark_task(
|
||||
state: &AppState,
|
||||
task_id: &str,
|
||||
status: &str,
|
||||
result: Option<&Value>,
|
||||
error: Option<&str>,
|
||||
) -> Result<(), AppError> {
|
||||
sqlx::query(
|
||||
"UPDATE tasks SET status = ?, result_json = ?, error = ?, finished_at = ? WHERE id = ?",
|
||||
)
|
||||
.bind(status)
|
||||
.bind(result.map(|v| v.to_string()))
|
||||
.bind(error)
|
||||
.bind(Utc::now().to_rfc3339())
|
||||
.bind(task_id)
|
||||
.execute(&state.pool)
|
||||
.await?;
|
||||
let level = if status == "success" { "info" } else { "error" };
|
||||
let message = match status {
|
||||
"success" => "任务成功完成",
|
||||
"timeout" => "任务执行超时",
|
||||
"cancelled" => "任务已取消",
|
||||
_ => "任务执行失败",
|
||||
};
|
||||
let data = serde_json::json!({
|
||||
"status": status,
|
||||
"error": error,
|
||||
"result": result
|
||||
});
|
||||
append_task_event(state, task_id, level, message, Some(&data)).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn audit(
|
||||
state: &AppState,
|
||||
user_id: Option<i64>,
|
||||
agent_id: Option<&str>,
|
||||
action: &str,
|
||||
target: Option<&str>,
|
||||
params_summary: Option<String>,
|
||||
success: bool,
|
||||
error: Option<&str>,
|
||||
) {
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO audit_logs (user_id, agent_id, action, target, params_summary, success, error) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(agent_id)
|
||||
.bind(action)
|
||||
.bind(target)
|
||||
.bind(params_summary)
|
||||
.bind(if success { 1 } else { 0 })
|
||||
.bind(error)
|
||||
.execute(&state.pool)
|
||||
.await;
|
||||
}
|
||||
Reference in New Issue
Block a user