alun_task/
worker.rs

1//! 任务消费者 —— 从 Kafka 消费任务并分发给 Handler 执行
2
3use std::sync::Arc;
4use std::time::Instant;
5use rdkafka::consumer::{StreamConsumer, Consumer};
6use rdkafka::ClientConfig;
7use rdkafka::message::BorrowedMessage;
8use rdkafka::producer::{FutureProducer, FutureRecord};
9use rdkafka::Message;
10use tokio::time::timeout;
11use tracing::{info, warn, error};
12
13use crate::storage::TaskStorage;
14use crate::HandlerRegistry;
15use crate::types::*;
16use crate::TaskMetrics;
17use crate::metrics::AtomicInc;
18
19/// 任务执行器
20///
21/// 从 Kafka 消费 `TaskMessage`，按 `task_type` 查找注册的 `TaskHandler`，
22/// 执行并委托 `TaskStorage` 记录结果、更新状态、处理重试和死信队列。
23/// 不持有任何 SQL 或表名——持久化完全交由 storage 代理。
24pub struct TaskWorker {
25    /// Kafka 消费者
26    consumer: Arc<StreamConsumer>,
27    /// 任务持久化接口
28    storage: Arc<dyn TaskStorage>,
29    /// 处理器注册中心
30    registry: HandlerRegistry,
31    /// 任务执行指标
32    metrics: Arc<TaskMetrics>,
33    /// TaskWorker 运行时配置
34    config: TaskWorkerConfig,
35    /// Kafka 生产者（用于发送死信消息）
36    producer: FutureProducer,
37    /// 运行状态标志
38    running: Arc<std::sync::atomic::AtomicBool>,
39}
40
41impl TaskWorker {
42    /// 创建任务执行器并订阅指定的 topics
43    ///
44    /// 注意：consumer 创建和 subscribe 必须在同一线程中完成，
45    /// 因此此方法在 spawn_blocking 中调用。
46    pub fn new(
47        config: TaskWorkerConfig,
48        storage: Arc<dyn TaskStorage>,
49        registry: HandlerRegistry,
50        metrics: Arc<TaskMetrics>,
51        topics: &[String],
52    ) -> Result<Self, String> {
53        let consumer: StreamConsumer = ClientConfig::new()
54            .set("bootstrap.servers", &config.brokers)
55            .set("group.id", &config.group_id)
56            .set("enable.auto.commit", "false")
57            .set("auto.offset.reset", "earliest")
58            .set("allow.auto.create.topics", "true")
59            .create()
60            .map_err(|e| format!("Kafka Consumer 创建失败: {}", e))?;
61
62        // 在同一线程中订阅 topics（rdkafka 要求 subscribe 与 consumer 创建在同一线程）
63        // 去重：同一 topic 注册多个 handler 时避免重复订阅
64        let mut unique_topics: Vec<&str> = topics.iter().map(|s| s.as_str()).collect();
65        unique_topics.sort_unstable();
66        unique_topics.dedup();
67        consumer
68            .subscribe(&unique_topics)
69            .map_err(|e| format!("Kafka Consumer 订阅失败: {}", e))?;
70
71        let producer: FutureProducer = ClientConfig::new()
72            .set("bootstrap.servers", &config.brokers)
73            .set("message.timeout.ms", "5000")
74            .create()
75            .map_err(|e| format!("Kafka DLQ Producer 创建失败: {}", e))?;
76
77        Ok(Self {
78            consumer: Arc::new(consumer),
79            storage,
80            registry,
81            metrics,
82            config,
83            producer,
84            running: Arc::new(std::sync::atomic::AtomicBool::new(true)),
85        })
86    }
87
88    /// 启动消费循环（subscribe 已在 new 中完成）
89    pub async fn run(&self) -> Result<(), String> {
90        info!("TaskWorker 启动，开始消费消息");
91
92        while self.running.load(std::sync::atomic::Ordering::Relaxed) {
93            match self.consumer.recv().await {
94                Err(e) => {
95                    error!("Kafka 接收消息失败: {}", e);
96                    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
97                }
98                Ok(msg) => {
99                    self.handle_message(&msg).await;
100                }
101            }
102        }
103
104        info!("TaskWorker 已停止");
105        Ok(())
106    }
107
108    /// 停止消费循环
109    pub fn stop(&self) {
110        self.running.store(false, std::sync::atomic::Ordering::Relaxed);
111    }
112
113    /// 处理单条 Kafka 消息：反序列化、查找 handler、执行、更新状态
114    async fn handle_message(&self, msg: &BorrowedMessage<'_>) {
115        let payload = match msg.payload() {
116            Some(p) => p,
117            None => {
118                warn!("收到空消息");
119                return;
120            }
121        };
122
123        let task_msg: TaskMessage = match serde_json::from_slice(payload) {
124            Ok(m) => m,
125            Err(e) => {
126                error!("消息反序列化失败: {}", e);
127                return;
128            }
129        };
130
131        if let Err(reason) = self.check_message_age(&task_msg) {
132            warn!(task_id = %task_msg.task_id, reason = reason, "消息已过期，跳过");
133            self.commit_offset(msg).await;
134            return;
135        }
136
137        self.metrics.total.inc();
138
139        let (handler, config) = match self.registry.get(task_msg.task_type) {
140            Some(h) => h,
141            None => {
142                warn!(task_type = task_msg.task_type, "未找到 handler");
143                self.commit_offset(msg).await;
144                return;
145            }
146        };
147
148        if let Err(e) = self.storage.update_task_status(&task_msg.task_id, TaskStatus::Processing).await {
149            error!(task_id = %task_msg.task_id, error = %e, "update_task_status(Processing) 失败");
150        }
151
152        let started_at = Instant::now();
153
154        let result = if config.timeout_seconds > 0 {
155            match timeout(
156                tokio::time::Duration::from_secs(config.timeout_seconds),
157                handler.execute(task_msg.payload.clone()),
158            )
159            .await
160            {
161                Ok(r) => r,
162                Err(_) => Err(format!("任务超时 ({}s)", config.timeout_seconds)),
163            }
164        } else {
165            handler.execute(task_msg.payload.clone()).await
166        };
167
168        let elapsed_ms = started_at.elapsed().as_millis() as i64;
169
170        match result {
171            Ok(output) => {
172                self.handle_success(&task_msg, &output, elapsed_ms).await;
173            }
174            Err(e) => {
175                self.handle_failure(&task_msg, &e, &config, elapsed_ms).await;
176            }
177        }
178
179        self.commit_offset(msg).await;
180    }
181
182    /// 异步提交 Kafka 消息 offset
183    async fn commit_offset(&self, msg: &BorrowedMessage<'_>) {
184        if let Err(e) = self.consumer.commit_message(msg, rdkafka::consumer::CommitMode::Async) {
185            error!(error = %e, "Kafka offset 提交失败");
186        }
187    }
188
189    /// 检查消息是否超过最大时效，超过则返回错误
190    fn check_message_age(&self, msg: &TaskMessage) -> Result<(), String> {
191        let submitted = chrono::DateTime::parse_from_rfc3339(&msg.submitted_at)
192            .map_err(|e| format!("解析 submitted_at 失败: {}", e))?;
193        let age = chrono::Utc::now()
194            .signed_duration_since(submitted.with_timezone(&chrono::Utc))
195            .num_seconds();
196        if age > self.config.max_message_age_secs as i64 {
197            return Err(format!(
198                "消息已超过最大时效 {}s（实际 {}s）",
199                self.config.max_message_age_secs, age
200            ));
201        }
202        Ok(())
203    }
204
205    /// 处理任务执行成功：更新状态为 Completed、存储结果、记录执行日志
206    async fn handle_success(&self, msg: &TaskMessage, output: &serde_json::Value, elapsed_ms: i64) {
207        if let Err(e) = self.storage.update_task_status(&msg.task_id, TaskStatus::Completed).await {
208            error!(task_id = %msg.task_id, error = %e, "update_task_status(Completed) 失败");
209        }
210        if let Err(e) = self.storage.save_task_result(&msg.task_id, output).await {
211            error!(task_id = %msg.task_id, error = %e, "save_task_result 失败");
212        }
213        if let Err(e) = self.storage.log_execution(&msg.task_id, TaskStatus::Completed, None, elapsed_ms).await {
214            error!(task_id = %msg.task_id, error = %e, "log_execution 失败");
215        }
216        self.metrics.completed.inc();
217        info!(task_id = %msg.task_id, elapsed_ms = elapsed_ms, "任务执行成功");
218    }
219
220    /// 处理任务执行失败：判断是否超过最大重试次数，决定是转入死信队列还是等待重试
221    async fn handle_failure(
222        &self,
223        msg: &TaskMessage,
224        err_msg: &str,
225        config: &TaskConfig,
226        elapsed_ms: i64,
227    ) {
228        self.metrics.failed.inc();
229
230        let current_retries = self.storage.get_retry_count(&msg.task_id).await.unwrap_or(0);
231        let attempt = current_retries + 1;
232
233        if attempt as u32 > config.max_retries {
234            if let Some(ref dlq_topic) = config.dead_letter_topic {
235                warn!(task_id = %msg.task_id, attempt = attempt, "转入死信队列");
236                if let Err(e) = self.send_to_dlq(msg, dlq_topic, err_msg).await {
237                    error!(task_id = %msg.task_id, error = %e, "send_to_dlq 失败");
238                }
239                if let Err(e) = self.storage.update_task_status(&msg.task_id, TaskStatus::DeadLetter).await {
240                    error!(task_id = %msg.task_id, error = %e, "update_task_status(DeadLetter) 失败");
241                }
242            } else {
243                warn!(task_id = %msg.task_id, attempt = attempt, "超过最大重试次数");
244                if let Err(e) = self.storage.update_task_status(&msg.task_id, TaskStatus::Failed).await {
245                    error!(task_id = %msg.task_id, error = %e, "update_task_status(Failed) 失败");
246                }
247                if let Err(e) = self.storage.save_task_result(
248                    &msg.task_id,
249                    &serde_json::json!({"error": err_msg, "retries": attempt}),
250                ).await {
251                    error!(task_id = %msg.task_id, error = %e, "save_task_result 失败");
252                }
253            }
254            if let Err(e) = self.storage.log_execution(&msg.task_id, TaskStatus::Failed, Some(err_msg), elapsed_ms).await {
255                error!(task_id = %msg.task_id, error = %e, "log_execution 失败");
256            }
257        } else {
258            if let Err(e) = self.storage.update_retry(&msg.task_id, attempt).await {
259                error!(task_id = %msg.task_id, error = %e, "update_retry 失败");
260            }
261            if let Err(e) = self.storage.log_execution(&msg.task_id, TaskStatus::Failed, Some(err_msg), elapsed_ms).await {
262                error!(task_id = %msg.task_id, error = %e, "log_execution 失败");
263            }
264            info!(task_id = %msg.task_id, attempt = attempt, "任务失败，等待重试: {}", err_msg);
265        }
266    }
267
268    /// 发送任务到死信队列（DLQ）
269    async fn send_to_dlq(
270        &self,
271        msg: &TaskMessage,
272        dead_letter_topic: &str,
273        reason: &str,
274    ) -> Result<(), String> {
275        let payload = serde_json::to_vec(msg).map_err(|e| format!("DLQ 序列化失败: {}", e))?;
276        let record = FutureRecord::to(dead_letter_topic)
277            .key(&msg.task_id)
278            .payload(&payload);
279
280        self.producer
281            .send(record, std::time::Duration::from_secs(5))
282            .await
283            .map_err(|(e, _)| format!("DLQ 发送失败: {}", e))?;
284
285        info!(task_id = %msg.task_id, reason = reason, "任务已转入死信队列: {}", dead_letter_topic);
286        Ok(())
287    }
288}
alun_task/worker.rs

alun_task/
worker.rs
