feat: implement full-stack TaskPilot scheduler
This commit is contained in:
@@ -0,0 +1,962 @@
|
||||
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,
|
||||
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<ExecutorInner>,
|
||||
}
|
||||
|
||||
struct ExecutorInner {
|
||||
pool: SqlitePool,
|
||||
logs_dir: PathBuf,
|
||||
shutdown: CancellationToken,
|
||||
running: Mutex<HashMap<String, RunningHandle>>,
|
||||
events: broadcast::Sender<ExecutionEvent>,
|
||||
http: reqwest::Client,
|
||||
log_config: LogRetentionConfig,
|
||||
cleanup_lock: AsyncMutex<()>,
|
||||
}
|
||||
|
||||
struct RunningHandle {
|
||||
run_id: Option<i64>,
|
||||
cancellation: CancellationToken,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct LogSink {
|
||||
file: Arc<AsyncMutex<tokio::fs::File>>,
|
||||
size: Arc<AtomicU64>,
|
||||
truncated: Arc<AtomicBool>,
|
||||
run_id: i64,
|
||||
events: broadcast::Sender<ExecutionEvent>,
|
||||
secrets: Arc<Vec<String>>,
|
||||
}
|
||||
|
||||
struct ScriptOutcome {
|
||||
status: RunStatus,
|
||||
exit_code: Option<i32>,
|
||||
error_message: Option<String>,
|
||||
}
|
||||
|
||||
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<TaskRunResponse> {
|
||||
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<ExecutionEvent> {
|
||||
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,
|
||||
interpreter,
|
||||
script,
|
||||
args,
|
||||
env,
|
||||
} => {
|
||||
let mut command = Command::new(interpreter);
|
||||
command
|
||||
.arg(script)
|
||||
.args(args)
|
||||
.envs(env)
|
||||
.current_dir(working_dir);
|
||||
self.run_command(task, run_id, cancellation, command).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).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,
|
||||
) -> 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}")),
|
||||
};
|
||||
|
||||
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<ExecutionEvent>,
|
||||
secrets: Vec<String>,
|
||||
) -> std::io::Result<Self> {
|
||||
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<HashMap<String, RunningHandle>>,
|
||||
) -> std::sync::MutexGuard<'_, HashMap<String, RunningHandle>> {
|
||||
running
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner())
|
||||
}
|
||||
|
||||
fn spawn_log_reader<R>(
|
||||
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<String> {
|
||||
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<String>) -> 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::<u64>();
|
||||
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::<u64>();
|
||||
assert!(total <= 10);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user