use std::{ collections::HashMap, path::{Path, PathBuf}, process::Stdio, sync::{ Arc, Mutex, atomic::{AtomicBool, AtomicU64, Ordering}, }, time::{Instant, SystemTime}, }; use chrono::Utc; use futures_util::StreamExt; use reqwest::Method; use serde::Serialize; use sqlx::SqlitePool; use tokio::{ io::{AsyncBufReadExt, AsyncRead, AsyncWriteExt, BufReader}, process::{Child, Command}, sync::{Mutex as AsyncMutex, broadcast}, time::{Duration, sleep, timeout}, }; use tokio_util::sync::CancellationToken; use crate::{ error::ApiResult, process_job::ProcessJob, python_env::prepare_script_environment, runs::{ model::{RunStatus, RunTrigger, TaskRunResponse}, repository, }, tasks::model::{HttpMethod, TaskAction, TaskResponse}, }; const MAX_LOG_BYTES: u64 = 5 * 1024 * 1024; const MAX_HTTP_LOG_BODY_BYTES: usize = 64 * 1024; const TERMINATION_WAIT_SECS: u64 = 5; const CREATE_NO_WINDOW: u32 = 0x0800_0000; const DEFAULT_RETENTION_DAYS: u64 = 30; const DEFAULT_MAX_TOTAL_LOG_BYTES: u64 = 512 * 1024 * 1024; #[derive(Clone, Debug)] pub struct LogRetentionConfig { pub retention_days: u64, pub max_total_bytes: u64, } impl Default for LogRetentionConfig { fn default() -> Self { Self { retention_days: DEFAULT_RETENTION_DAYS, max_total_bytes: DEFAULT_MAX_TOTAL_LOG_BYTES, } } } #[derive(Clone, Debug, Serialize)] #[serde(tag = "event", content = "data", rename_all = "snake_case")] pub enum ExecutionEvent { Run { run: TaskRunResponse, }, Log { run_id: i64, offset: u64, next_offset: u64, content: String, }, } impl ExecutionEvent { pub const fn event_name(&self) -> &'static str { match self { Self::Run { .. } => "run", Self::Log { .. } => "log", } } } #[derive(Clone)] pub struct TaskExecutor { inner: Arc, } struct ExecutorInner { pool: SqlitePool, logs_dir: PathBuf, shutdown: CancellationToken, running: Mutex>, events: broadcast::Sender, http: reqwest::Client, log_config: LogRetentionConfig, cleanup_lock: AsyncMutex<()>, } struct RunningHandle { run_id: Option, cancellation: CancellationToken, } #[derive(Clone)] struct LogSink { file: Arc>, size: Arc, truncated: Arc, run_id: i64, events: broadcast::Sender, secrets: Arc>, } struct ScriptOutcome { status: RunStatus, exit_code: Option, error_message: Option, } enum ProcessEnd { Exited(std::process::ExitStatus), TimedOut, Cancelled, } impl TaskExecutor { pub fn new(pool: SqlitePool, logs_dir: PathBuf, log_config: LogRetentionConfig) -> Self { let (events, _) = broadcast::channel(512); Self { inner: Arc::new(ExecutorInner { pool, logs_dir, shutdown: CancellationToken::new(), running: Mutex::new(HashMap::new()), events, http: reqwest::Client::builder() .user_agent("TaskPilot/0.1") .redirect(reqwest::redirect::Policy::limited(10)) .build() .unwrap_or_default(), log_config, cleanup_lock: AsyncMutex::new(()), }), } } pub async fn trigger( &self, task: &TaskResponse, trigger: RunTrigger, ) -> ApiResult { let cancellation = CancellationToken::new(); let already_running = { let mut running = lock_running(&self.inner.running); if running.contains_key(&task.id) { true } else { running.insert( task.id.clone(), RunningHandle { run_id: None, cancellation: cancellation.clone(), }, ); false } }; if already_running { let run = repository::create_skipped( &self.inner.pool, &task.id, trigger, "同一任务已有运行实例,本次触发已跳过", ) .await?; self.emit_run(&run); return Ok(run); } let run = match repository::create_queued(&self.inner.pool, &task.id, trigger).await { Ok(run) => run, Err(error) => { lock_running(&self.inner.running).remove(&task.id); return Err(error); } }; if let Some(handle) = lock_running(&self.inner.running).get_mut(&task.id) { handle.run_id = Some(run.id); } self.emit_run(&run); let executor = self.clone(); let task = task.clone(); tokio::spawn(async move { executor.execute_run(task, run.id, cancellation).await; }); Ok(run) } pub fn cancel_task(&self, task_id: &str) -> bool { let mut running = lock_running(&self.inner.running); if let Some(handle) = running.remove(task_id) { handle.cancellation.cancel(); true } else { false } } pub fn cancel_run(&self, run_id: i64) -> bool { let mut running = lock_running(&self.inner.running); let task_id = running .iter() .find_map(|(task_id, handle)| (handle.run_id == Some(run_id)).then(|| task_id.clone())); if let Some(task_id) = task_id && let Some(handle) = running.remove(&task_id) { handle.cancellation.cancel(); return true; } false } pub fn subscribe(&self) -> broadcast::Receiver { self.inner.events.subscribe() } pub fn log_path(&self, task_id: &str, run_id: i64) -> PathBuf { self.inner .logs_dir .join(task_id) .join(format!("{run_id}.log")) } pub fn shutdown(&self) { self.inner.shutdown.cancel(); let running = lock_running(&self.inner.running); for handle in running.values() { handle.cancellation.cancel(); } } pub async fn cleanup_logs(&self) -> std::io::Result<()> { let _guard = self.inner.cleanup_lock.lock().await; let logs_dir = self.inner.logs_dir.clone(); let config = self.inner.log_config.clone(); tokio::task::spawn_blocking(move || cleanup_logs_sync(&logs_dir, &config)) .await .map_err(std::io::Error::other)? } async fn execute_run(&self, task: TaskResponse, run_id: i64, cancellation: CancellationToken) { let started_at = Utc::now(); let started = Instant::now(); if let Err(error) = repository::mark_running(&self.inner.pool, run_id, started_at).await { tracing::error!(run_id, error = %error, "failed to mark task run as running"); lock_running(&self.inner.running).remove(&task.id); return; } self.emit_run_by_id(run_id).await; let outcome = self.run_action(&task, run_id, &cancellation).await; let finished_at = Utc::now(); let duration_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX); if let Err(error) = repository::complete( &self.inner.pool, run_id, outcome.status, finished_at, duration_ms, outcome.exit_code, outcome.error_message.as_deref(), ) .await { tracing::error!(run_id, error = %error, "failed to complete task run"); } self.emit_run_by_id(run_id).await; lock_running(&self.inner.running).remove(&task.id); if let Err(error) = self.cleanup_logs().await { tracing::warn!(%error, "failed to clean task logs"); } } async fn run_action( &self, task: &TaskResponse, run_id: i64, cancellation: &CancellationToken, ) -> ScriptOutcome { match &task.action { TaskAction::Script { working_dir, venv, interpreter, script, args, env, } => { let prepared = match prepare_script_environment( working_dir, interpreter, venv.as_deref(), env, ) { Ok(prepared) => prepared, Err(error) => { return failed_outcome(format!("准备 Python 运行环境失败: {error}")); } }; let mut command = Command::new(&prepared.interpreter); prepared.apply(&mut command); command.arg(script).args(args).current_dir(working_dir); let startup_note = prepared .virtual_env .as_ref() .map(|venv| format!("使用 Python 虚拟环境: {}", venv.display())); self.run_command(task, run_id, cancellation, command, startup_note) .await } TaskAction::Process { executable, working_dir, args, env, } => { let mut command = Command::new(executable); command.args(args).envs(env); if let Some(working_dir) = working_dir { command.current_dir(working_dir); } self.run_command(task, run_id, cancellation, command, None) .await } TaskAction::Http { .. } => self.run_http(task, run_id, cancellation).await, } } async fn run_command( &self, task: &TaskResponse, run_id: i64, cancellation: &CancellationToken, mut command: Command, startup_note: Option, ) -> ScriptOutcome { let log_path = self.log_path(&task.id, run_id); let secrets = sensitive_values(&task.action); let sink = match LogSink::create( &log_path, task, run_id, self.inner.events.clone(), secrets, ) .await { Ok(sink) => sink, Err(error) => return failed_outcome(format!("无法创建运行日志: {error}")), }; if let Some(note) = startup_note { sink.system(¬e).await; } let job = match ProcessJob::new() { Ok(job) => job, Err(error) => { sink.system(&format!("无法创建 Windows Job Object: {error}")) .await; return failed_outcome(format!("无法创建 Windows Job Object: {error}")); } }; command .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .kill_on_drop(true); #[cfg(windows)] { command.creation_flags(CREATE_NO_WINDOW); } let mut child = match command.spawn() { Ok(child) => child, Err(error) => { sink.system(&format!("启动进程失败: {error}")).await; return failed_outcome(format!("启动进程失败: {error}")); } }; if let Err(error) = job.assign(&child) { let _ = child.kill().await; sink.system(&format!("关联 Job Object 失败: {error}")).await; return failed_outcome(format!("关联 Windows Job Object 失败: {error}")); } let stdout_task = child .stdout .take() .map(|stdout| spawn_log_reader(stdout, "stdout", sink.clone())); let stderr_task = child .stderr .take() .map(|stderr| spawn_log_reader(stderr, "stderr", sink.clone())); let process_end = tokio::select! { status = child.wait() => { match status { Ok(status) => ProcessEnd::Exited(status), Err(error) => { sink.system(&format!("等待进程失败: {error}")).await; return failed_outcome(format!("等待进程失败: {error}")); } } } _ = sleep(Duration::from_secs(task.timeout_secs)) => { if let Err(error) = job.terminate() { tracing::error!(task_id = %task.id, run_id, %error, "failed to terminate timed out job"); } wait_after_termination(&mut child).await; ProcessEnd::TimedOut } _ = cancellation.cancelled() => { if let Err(error) = job.terminate() { tracing::error!(task_id = %task.id, run_id, %error, "failed to terminate cancelled job"); } wait_after_termination(&mut child).await; ProcessEnd::Cancelled } _ = self.inner.shutdown.cancelled() => { if let Err(error) = job.terminate() { tracing::error!(task_id = %task.id, run_id, %error, "failed to terminate job during shutdown"); } wait_after_termination(&mut child).await; ProcessEnd::Cancelled } }; if let Some(stdout_task) = stdout_task && let Err(error) = stdout_task.await { tracing::warn!(task_id = %task.id, run_id, %error, "stdout capture task failed"); } if let Some(stderr_task) = stderr_task && let Err(error) = stderr_task.await { tracing::warn!(task_id = %task.id, run_id, %error, "stderr capture task failed"); } match process_end { ProcessEnd::Exited(status) => { let exit_code = status.code(); if status.success() { sink.system("任务执行成功").await; ScriptOutcome { status: RunStatus::Succeeded, exit_code, error_message: None, } } else { let message = format!( "进程退出码为 {}", exit_code.map_or_else(|| "unknown".to_owned(), |code| code.to_string()) ); sink.system(&message).await; ScriptOutcome { status: RunStatus::Failed, exit_code, error_message: Some(message), } } } ProcessEnd::TimedOut => { let message = format!("任务执行超过 {} 秒,已终止进程树", task.timeout_secs); sink.system(&message).await; ScriptOutcome { status: RunStatus::TimedOut, exit_code: None, error_message: Some(message), } } ProcessEnd::Cancelled => { let message = "任务已取消".to_owned(); sink.system(&message).await; ScriptOutcome { status: RunStatus::Cancelled, exit_code: None, error_message: Some(message), } } } } async fn run_http( &self, task: &TaskResponse, run_id: i64, cancellation: &CancellationToken, ) -> ScriptOutcome { let TaskAction::Http { url, method, headers, body, expected_status, } = &task.action else { return failed_outcome("invalid HTTP action payload"); }; let log_path = self.log_path(&task.id, run_id); let sink = match LogSink::create( &log_path, task, run_id, self.inner.events.clone(), sensitive_values(&task.action), ) .await { Ok(sink) => sink, Err(error) => return failed_outcome(format!("无法创建运行日志: {error}")), }; let method = match method { HttpMethod::Get => Method::GET, HttpMethod::Post => Method::POST, }; let mut request = self .inner .http .request(method.clone(), url) .timeout(Duration::from_secs(task.timeout_secs)); for (name, value) in headers { request = request.header(name, value); } if let Some(body) = body { request = request.body(body.clone()); } sink.http(&format!("请求 {method} {url}")).await; let response = tokio::select! { response = request.send() => { match response { Ok(response) => response, Err(error) => { let message = if error.is_timeout() { format!("HTTP 请求超过 {} 秒", task.timeout_secs) } else { format!("HTTP 请求失败: {error}") }; sink.system(&message).await; return ScriptOutcome { status: if error.is_timeout() { RunStatus::TimedOut } else { RunStatus::Failed }, exit_code: None, error_message: Some(message), }; } } } _ = cancellation.cancelled() => { sink.system("HTTP 请求已取消").await; return ScriptOutcome { status: RunStatus::Cancelled, exit_code: None, error_message: Some("HTTP 请求已取消".to_owned()), }; } _ = self.inner.shutdown.cancelled() => { sink.system("HTTP 请求因服务关闭而取消").await; return ScriptOutcome { status: RunStatus::Cancelled, exit_code: None, error_message: Some("HTTP 请求因服务关闭而取消".to_owned()), }; } }; let http_status = response.status(); sink.http(&format!("响应状态 {http_status}")).await; let mut response_stream = response.bytes_stream(); let mut response_body = Vec::new(); let mut body_truncated = false; loop { let next = tokio::select! { chunk = response_stream.next() => chunk, _ = cancellation.cancelled() => { sink.system("读取 HTTP 响应已取消").await; return ScriptOutcome { status: RunStatus::Cancelled, exit_code: None, error_message: Some("读取 HTTP 响应已取消".to_owned()), }; } _ = self.inner.shutdown.cancelled() => { sink.system("HTTP 响应读取因服务关闭而取消").await; return ScriptOutcome { status: RunStatus::Cancelled, exit_code: None, error_message: Some("HTTP 响应读取因服务关闭而取消".to_owned()), }; } }; match next { Some(Ok(chunk)) => { let remaining = MAX_HTTP_LOG_BODY_BYTES.saturating_sub(response_body.len()); let copy_length = remaining.min(chunk.len()); response_body.extend_from_slice(&chunk[..copy_length]); if chunk.len() > copy_length { body_truncated = true; break; } } Some(Err(error)) => { let message = if error.is_timeout() { format!("读取 HTTP 响应超过 {} 秒", task.timeout_secs) } else { format!("读取 HTTP 响应失败: {error}") }; sink.system(&message).await; return ScriptOutcome { status: if error.is_timeout() { RunStatus::TimedOut } else { RunStatus::Failed }, exit_code: None, error_message: Some(message), }; } None => break, } } let body_text = String::from_utf8_lossy(&response_body); if !body_text.is_empty() { sink.http(&body_text).await; } if body_truncated { sink.system(&format!( "HTTP 响应体超过 {MAX_HTTP_LOG_BODY_BYTES} 字节,日志仅保留前部内容" )) .await; } let expected = if expected_status.is_empty() { http_status.is_success() } else { expected_status.contains(&http_status.as_u16()) }; if expected { sink.system("HTTP 任务执行成功").await; ScriptOutcome { status: RunStatus::Succeeded, exit_code: Some(i32::from(http_status.as_u16())), error_message: None, } } else { let message = format!("HTTP 状态码 {http_status} 不在期望范围内"); sink.system(&message).await; ScriptOutcome { status: RunStatus::Failed, exit_code: Some(i32::from(http_status.as_u16())), error_message: Some(message), } } } fn emit_run(&self, run: &TaskRunResponse) { let _ = self .inner .events .send(ExecutionEvent::Run { run: run.clone() }); } async fn emit_run_by_id(&self, run_id: i64) { match repository::get(&self.inner.pool, run_id).await { Ok(Some(run)) => self.emit_run(&run), Ok(None) => {} Err(error) => { tracing::warn!(run_id, error = %error, "failed to load run for SSE event"); } } } } impl LogSink { async fn create( path: &Path, task: &TaskResponse, run_id: i64, events: broadcast::Sender, secrets: Vec, ) -> std::io::Result { if let Some(parent) = path.parent() { tokio::fs::create_dir_all(parent).await?; } let header = format!( "[system] TaskPilot run {run_id} for task {} started at {}\n", task.id, Utc::now().to_rfc3339() ); tokio::fs::write(path, header.as_bytes()).await?; let file = tokio::fs::OpenOptions::new() .append(true) .open(path) .await?; let mut secrets = secrets; secrets.sort_by_key(|value| std::cmp::Reverse(value.len())); Ok(Self { file: Arc::new(AsyncMutex::new(file)), size: Arc::new(AtomicU64::new(header.len() as u64)), truncated: Arc::new(AtomicBool::new(false)), run_id, events, secrets: Arc::new(secrets), }) } async fn append(&self, text: &str) { let redacted = redact_text(text, &self.secrets); let bytes = redacted.as_bytes(); let current = self.size.load(Ordering::Relaxed); if current + bytes.len() as u64 > MAX_LOG_BYTES { if !self.truncated.swap(true, Ordering::Relaxed) { let notice = format!("[system] 日志超过 {MAX_LOG_BYTES} 字节,后续输出已丢弃\n"); let mut file = self.file.lock().await; let _ = file.write_all(notice.as_bytes()).await; } return; } let mut file = self.file.lock().await; if file.write_all(bytes).await.is_err() { return; } let next = self.size.fetch_add(bytes.len() as u64, Ordering::Relaxed) + bytes.len() as u64; let _ = self.events.send(ExecutionEvent::Log { run_id: self.run_id, offset: current, next_offset: next, content: redacted, }); } async fn system(&self, message: &str) { self.append(&format!("[system] {message}\n")).await; } async fn http(&self, message: &str) { self.append(&format!("[http] {message}\n")).await; } } fn lock_running( running: &Mutex>, ) -> std::sync::MutexGuard<'_, HashMap> { running .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) } fn spawn_log_reader( reader: R, stream: &'static str, sink: LogSink, ) -> tokio::task::JoinHandle<()> where R: AsyncRead + Unpin + Send + 'static, { tokio::spawn(async move { let mut reader = BufReader::new(reader); let mut line = Vec::new(); loop { line.clear(); match reader.read_until(b'\n', &mut line).await { Ok(0) => break, Ok(_) => { let text = String::from_utf8_lossy(&line); sink.append(&format!("[{stream}] {text}")).await; } Err(error) => { tracing::warn!(stream, %error, "failed to read task output"); break; } } } }) } async fn wait_after_termination(child: &mut Child) { if timeout(Duration::from_secs(TERMINATION_WAIT_SECS), child.wait()) .await .is_err() { let _ = child.kill().await; } } fn sensitive_values(action: &TaskAction) -> Vec { match action { TaskAction::Script { env, .. } | TaskAction::Process { env, .. } => env .iter() .filter(|(key, value)| is_sensitive_name(key) && !value.is_empty()) .map(|(_, value)| value.clone()) .collect(), TaskAction::Http { headers, .. } => headers .iter() .filter(|(key, value)| is_sensitive_name(key) && !value.is_empty()) .map(|(_, value)| value.clone()) .collect(), } } fn is_sensitive_name(name: &str) -> bool { let name = name.to_ascii_uppercase(); [ "SECRET", "TOKEN", "PASSWORD", "PASS", "API_KEY", "PRIVATE_KEY", "CREDENTIAL", "AUTHORIZATION", "AUTH", ] .iter() .any(|marker| name.contains(marker)) } fn redact_text(text: &str, secrets: &[String]) -> String { secrets.iter().fold(text.to_owned(), |current, secret| { current.replace(secret, "[REDACTED]") }) } fn failed_outcome(message: impl Into) -> ScriptOutcome { ScriptOutcome { status: RunStatus::Failed, exit_code: None, error_message: Some(message.into()), } } fn cleanup_logs_sync(root: &Path, config: &LogRetentionConfig) -> std::io::Result<()> { if !root.exists() { return Ok(()); } let mut files = Vec::new(); collect_log_files(root, &mut files)?; let retention = Duration::from_secs(config.retention_days.saturating_mul(86_400)); let cutoff = SystemTime::now().checked_sub(retention); if let Some(cutoff) = cutoff { files.retain(|(path, metadata)| { let expired = metadata .modified() .map(|modified| modified < cutoff) .unwrap_or(false); if expired { let _ = std::fs::remove_file(path); false } else { true } }); } files.sort_by_key(|(_, metadata)| metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH)); let mut total = files .iter() .map(|(_, metadata)| metadata.len()) .sum::(); for (path, metadata) in files { if total <= config.max_total_bytes { break; } if std::fs::remove_file(&path).is_ok() { total = total.saturating_sub(metadata.len()); } } remove_empty_directories(root)?; Ok(()) } fn collect_log_files( directory: &Path, files: &mut Vec<(PathBuf, std::fs::Metadata)>, ) -> std::io::Result<()> { for entry in std::fs::read_dir(directory)? { let entry = entry?; let file_type = entry.file_type()?; if file_type.is_symlink() { continue; } if file_type.is_dir() { collect_log_files(&entry.path(), files)?; } else if file_type.is_file() { files.push((entry.path(), entry.metadata()?)); } } Ok(()) } fn remove_empty_directories(root: &Path) -> std::io::Result<()> { for entry in std::fs::read_dir(root)? { let entry = entry?; if entry.file_type()?.is_dir() { let path = entry.path(); remove_empty_directories(&path)?; if std::fs::read_dir(&path)?.next().is_none() { let _ = std::fs::remove_dir(path); } } } Ok(()) } #[cfg(test)] mod tests { use super::*; #[test] fn terminal_statuses_are_identified() { assert!(!RunStatus::Queued.is_terminal()); assert!(!RunStatus::Running.is_terminal()); assert!(RunStatus::Succeeded.is_terminal()); assert!(RunStatus::Failed.is_terminal()); assert!(RunStatus::TimedOut.is_terminal()); assert!(RunStatus::Cancelled.is_terminal()); assert!(RunStatus::Skipped.is_terminal()); } #[test] fn redacts_sensitive_values() { let secrets = vec!["super-secret".to_owned(), "token-value".to_owned()]; let output = redact_text("found super-secret and token-value", &secrets); assert_eq!(output, "found [REDACTED] and [REDACTED]"); } #[test] fn identifies_sensitive_environment_names() { assert!(is_sensitive_name("API_TOKEN")); assert!(is_sensitive_name("DB_PASSWORD")); assert!(is_sensitive_name("Authorization")); assert!(!is_sensitive_name("PYTHONUNBUFFERED")); } #[test] fn cleanup_enforces_total_log_capacity() { let temp = tempfile::TempDir::new().unwrap(); for index in 0..6 { std::fs::write(temp.path().join(format!("run-{index}.log")), vec![b'x'; 4]).unwrap(); std::thread::sleep(std::time::Duration::from_millis(5)); } cleanup_logs_sync( temp.path(), &LogRetentionConfig { retention_days: 30, max_total_bytes: 10, }, ) .unwrap(); let total = std::fs::read_dir(temp.path()) .unwrap() .map(|entry| entry.unwrap().metadata().unwrap().len()) .sum::(); assert!(total <= 10); } }