feat: 拆分 worker 架构并增强 assistant 安全保护
- 新增 daemon/channel/assistant/function/scheduler/logging worker 架构和 JetStream 消息协议 - 修复邮件通知误入 assistant、迁移 checksum、日志 subject 和工具调用无限循环风险 - 新增 assistant 工具轮次熔断、工具结果截断和提示词软收敛规则 - 补充邮件推送关键字、可信发件人配置、升级日志和验证测试
This commit is contained in:
@@ -0,0 +1,18 @@
|
||||
[package]
|
||||
name = "ias-scheduler"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
ias-common = { path = "../ias-common" }
|
||||
anyhow = "1"
|
||||
chrono = { version = "0.4", features = ["serde"] }
|
||||
cron = "0.12"
|
||||
serde_json = "1"
|
||||
sqlx = { version = "0.9.0", features = [
|
||||
"runtime-tokio",
|
||||
"tls-rustls",
|
||||
"postgres",
|
||||
"chrono",
|
||||
] }
|
||||
tokio = { version = "1.0", features = ["rt-multi-thread", "macros", "time"] }
|
||||
@@ -0,0 +1,265 @@
|
||||
use std::str::FromStr;
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use cron::Schedule;
|
||||
use ias_common::mq::{
|
||||
JetStreamBus, JetStreamConfig, MessageEnvelope, ScheduledTriggerPayload,
|
||||
schedule_trigger_subject,
|
||||
};
|
||||
use serde_json::Value;
|
||||
use sqlx::PgPool;
|
||||
|
||||
const DEFAULT_FIRE_RECLAIM_TIMEOUT_SECS: u64 = 120;
|
||||
|
||||
#[derive(Debug, sqlx::FromRow)]
|
||||
struct DueSchedule {
|
||||
id: i64,
|
||||
schedule_type: String,
|
||||
expression: String,
|
||||
target_topic: String,
|
||||
payload: Value,
|
||||
next_fire_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
pub async fn run_worker(pool: PgPool) -> anyhow::Result<()> {
|
||||
let config = JetStreamConfig::from_env();
|
||||
loop {
|
||||
match run_once(pool.clone(), config.clone()).await {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(error) => {
|
||||
eprintln!("Scheduler Worker 连接失败: {error:#}");
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_once(pool: PgPool, config: JetStreamConfig) -> anyhow::Result<()> {
|
||||
let bus = JetStreamBus::connect(config).await?;
|
||||
bus.ensure_streams().await?;
|
||||
let tick_secs = read_u64_env("IAS_SCHEDULER_TICK_SECS", 10);
|
||||
let fire_reclaim_timeout_secs = read_u64_env(
|
||||
"IAS_SCHEDULER_FIRE_RECLAIM_TIMEOUT_SECS",
|
||||
DEFAULT_FIRE_RECLAIM_TIMEOUT_SECS,
|
||||
);
|
||||
let mut tick = tokio::time::interval(Duration::from_secs(tick_secs));
|
||||
|
||||
loop {
|
||||
tick.tick().await;
|
||||
let schedules = load_due_schedules(&pool).await?;
|
||||
for schedule in schedules {
|
||||
if let Err(error) =
|
||||
publish_schedule(&pool, &bus, schedule, fire_reclaim_timeout_secs).await
|
||||
{
|
||||
eprintln!("发布定时任务失败: {error:#}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_due_schedules(pool: &PgPool) -> anyhow::Result<Vec<DueSchedule>> {
|
||||
let schedules = sqlx::query_as::<_, DueSchedule>(
|
||||
r#"
|
||||
SELECT id, schedule_type, expression, target_topic, payload, next_fire_at
|
||||
FROM ias_schedules
|
||||
WHERE enabled = TRUE
|
||||
AND next_fire_at <= NOW()
|
||||
ORDER BY next_fire_at ASC
|
||||
LIMIT 100
|
||||
"#,
|
||||
)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
Ok(schedules)
|
||||
}
|
||||
|
||||
async fn publish_schedule(
|
||||
pool: &PgPool,
|
||||
bus: &JetStreamBus,
|
||||
schedule: DueSchedule,
|
||||
fire_reclaim_timeout_secs: u64,
|
||||
) -> anyhow::Result<()> {
|
||||
let fire_id = schedule_fire_id(schedule.id, schedule.next_fire_at);
|
||||
let Some(attempt_count) =
|
||||
claim_fire(pool, &fire_id, &schedule, fire_reclaim_timeout_secs).await?
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let payload = ScheduledTriggerPayload {
|
||||
fire_id: fire_id.clone(),
|
||||
schedule_id: schedule.id,
|
||||
schedule_type: schedule.schedule_type.clone(),
|
||||
target_topic: schedule.target_topic.clone(),
|
||||
fire_time_ms: schedule.next_fire_at.timestamp_millis(),
|
||||
payload: schedule.payload,
|
||||
};
|
||||
let subject = schedule_trigger_subject(&schedule.target_topic);
|
||||
let envelope = MessageEnvelope::new(
|
||||
subject.as_str(),
|
||||
format!("schedule/{fire_id}"),
|
||||
format!("schedule-{fire_id}"),
|
||||
payload,
|
||||
)
|
||||
.with_attempt_hint(attempt_count.saturating_sub(1) as u32);
|
||||
match bus.publish_envelope(subject.as_str(), &envelope).await {
|
||||
Ok(()) => {
|
||||
let next_fire_at = next_fire_at(
|
||||
&schedule.schedule_type,
|
||||
&schedule.expression,
|
||||
schedule.next_fire_at,
|
||||
)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE ias_schedules
|
||||
SET last_fired_at = $1, next_fire_at = $2, updated_at = NOW()
|
||||
WHERE id = $3
|
||||
"#,
|
||||
)
|
||||
.bind(schedule.next_fire_at)
|
||||
.bind(next_fire_at)
|
||||
.bind(schedule.id)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE ias_schedule_fires
|
||||
SET status = 'published', published_at = NOW(), updated_at = NOW()
|
||||
WHERE fire_id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(&fire_id)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
}
|
||||
Err(error) => {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE ias_schedule_fires
|
||||
SET status = 'failed', error = $2, updated_at = NOW()
|
||||
WHERE fire_id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(&fire_id)
|
||||
.bind(error.to_string())
|
||||
.execute(pool)
|
||||
.await?;
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn claim_fire(
|
||||
pool: &PgPool,
|
||||
fire_id: &str,
|
||||
schedule: &DueSchedule,
|
||||
fire_reclaim_timeout_secs: u64,
|
||||
) -> anyhow::Result<Option<i32>> {
|
||||
let timeout_secs = fire_reclaim_timeout_secs.min(i32::MAX as u64) as i32;
|
||||
let attempt_count = sqlx::query_scalar::<_, i32>(
|
||||
r#"
|
||||
INSERT INTO ias_schedule_fires (
|
||||
fire_id,
|
||||
schedule_id,
|
||||
fire_time,
|
||||
status,
|
||||
attempt_count,
|
||||
error,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, 'publishing', 1, NULL, NOW())
|
||||
ON CONFLICT (fire_id) DO UPDATE SET
|
||||
status = 'publishing',
|
||||
attempt_count = ias_schedule_fires.attempt_count + 1,
|
||||
error = NULL,
|
||||
updated_at = NOW()
|
||||
WHERE ias_schedule_fires.status = 'failed'
|
||||
OR (
|
||||
ias_schedule_fires.status = 'publishing'
|
||||
AND ias_schedule_fires.updated_at < NOW() - ($4::INT * INTERVAL '1 second')
|
||||
)
|
||||
RETURNING attempt_count
|
||||
"#,
|
||||
)
|
||||
.bind(fire_id)
|
||||
.bind(schedule.id)
|
||||
.bind(schedule.next_fire_at)
|
||||
.bind(timeout_secs)
|
||||
.fetch_optional(pool)
|
||||
.await?;
|
||||
Ok(attempt_count)
|
||||
}
|
||||
|
||||
fn schedule_fire_id(schedule_id: i64, fire_time: DateTime<Utc>) -> String {
|
||||
format!("{}/{}", schedule_id, fire_time.timestamp_millis())
|
||||
}
|
||||
|
||||
fn next_fire_at(
|
||||
schedule_type: &str,
|
||||
expression: &str,
|
||||
current_fire_at: DateTime<Utc>,
|
||||
) -> anyhow::Result<DateTime<Utc>> {
|
||||
match schedule_type {
|
||||
"interval" => {
|
||||
let seconds = expression.trim().parse::<i64>().unwrap_or(60).max(1);
|
||||
Ok(current_fire_at + chrono::Duration::seconds(seconds))
|
||||
}
|
||||
"cron" => {
|
||||
let schedule = Schedule::from_str(expression)?;
|
||||
schedule
|
||||
.after(¤t_fire_at)
|
||||
.next()
|
||||
.ok_or_else(|| anyhow::anyhow!("cron 表达式没有下一次触发时间"))
|
||||
}
|
||||
other => anyhow::bail!("不支持的 schedule_type: {other}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn read_u64_env(name: &str, default: u64) -> u64 {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.and_then(|value| value.trim().parse::<u64>().ok())
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn interval_next_fire_at_advances_by_expression_seconds() {
|
||||
let current = DateTime::parse_from_rfc3339("2026-07-09T00:00:00Z")
|
||||
.unwrap()
|
||||
.with_timezone(&Utc);
|
||||
|
||||
assert_eq!(
|
||||
next_fire_at("interval", "30", current).unwrap(),
|
||||
current + chrono::Duration::seconds(30)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_next_fire_at_uses_cron_expression() {
|
||||
let current = DateTime::parse_from_rfc3339("2026-07-09T00:00:00Z")
|
||||
.unwrap()
|
||||
.with_timezone(&Utc);
|
||||
|
||||
assert_eq!(
|
||||
next_fire_at("cron", "0 0/5 * * * * *", current).unwrap(),
|
||||
current + chrono::Duration::minutes(5)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schedule_fire_id_is_stable_for_schedule_and_fire_time() {
|
||||
let current = DateTime::parse_from_rfc3339("2026-07-09T00:00:00Z")
|
||||
.unwrap()
|
||||
.with_timezone(&Utc);
|
||||
|
||||
assert_eq!(schedule_fire_id(42, current), "42/1783555200000");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user