sz-rust-core 0.6.5

SZ-Rust 核心库:HTTP 服务器、路由、控制器、中间件,对标 ThinkPHP 8
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
//! sz-orm-queue 消费者接入
//!
//! ## PHP 对齐
//!
//! 对齐 PHP `think-queue` 的消费者模型:
//!
//! ```php
//! // think-queue 的 worker 模型
//! $worker->runNextJob($connection, $queue, $delay, $sleep, $tries);
//! ```
//!
//! Rust 端复用 `sz-orm-queue::MessageQueue` trait,提供 consumer loop helper:
//!
//! - `QueueConsumer` trait:业务侧实现消息处理逻辑
//! - `QueueRuntime`:管理 consumer lifecycle(启动 / 停止 / 监听 cancel)
//!
//! ## 设计
//!
//! sz-orm-queue 的 `MessageQueue` trait 只提供原子操作(publish/consume/ack),
//! 没有 consumer loop helper。本模块补充该能力,配合 `CancellationToken` 做退出。

use std::sync::Arc;
use std::time::Duration;

use async_trait::async_trait;
use tokio_util::sync::CancellationToken;

use crate::orm::MessageQueue;

/// 队列消费者错误
#[derive(Debug, Clone, thiserror::Error)]
pub enum QueueConsumerError {
    /// 消息处理失败(消息会被 nack,留在 in_flight 中)
    #[error("consumer error: {0}")]
    Handler(String),
    /// 底层队列操作失败
    #[error("queue error: {0}")]
    Queue(String),
}

/// 队列消费者 trait
///
/// 业务侧实现该 trait,处理从队列消费的消息。
///
/// ## 行为契约
///
/// - 返回 `Ok(())`:自动调用 `queue.ack(message_id)` 确认消息
/// - 返回 `Err(_)`:不 ack(消息留在 `in_flight`,需人工干预或超时重投)
#[async_trait]
pub trait QueueConsumer: Send + Sync {
    /// 处理一条消息
    ///
    /// - 成功返回 `Ok(())` 会触发自动 ack
    /// - 失败返回 `Err(_)` 会跳过 ack(消息留在 in_flight)
    async fn handle(&self, message: &sz_orm_queue::Message) -> Result<(), QueueConsumerError>;
}

/// 队列运行时配置
#[derive(Debug, Clone)]
pub struct QueueRuntimeConfig {
    /// 消费主题
    pub topic: String,
    /// 队列为空时的轮询间隔(毫秒)
    pub poll_interval_ms: u64,
    /// 单条消息处理最大重试次数(0 表示不重试)
    pub max_retries: u32,
}

impl Default for QueueRuntimeConfig {
    fn default() -> Self {
        Self {
            topic: "default".to_string(),
            poll_interval_ms: 100,
            max_retries: 0,
        }
    }
}

impl QueueRuntimeConfig {
    /// 创建新配置
    pub fn new(topic: impl Into<String>) -> Self {
        Self {
            topic: topic.into(),
            ..Default::default()
        }
    }

    /// 自定义轮询间隔
    pub fn with_poll_interval(mut self, ms: u64) -> Self {
        self.poll_interval_ms = ms;
        self
    }

    /// 自定义最大重试次数
    pub fn with_max_retries(mut self, n: u32) -> Self {
        self.max_retries = n;
        self
    }
}

/// 队列运行时
///
/// 管理 consumer lifecycle:启动一个消费循环任务,监听 `CancellationToken` 优雅退出。
///
/// ## 设计
///
/// - **不持有 `JoinHandle`**:消费任务由调用方持有,本结构仅提供启动入口
/// - **ack 策略**:handler 返回 Ok 时自动 ack;Err 时不 ack(消息留在 in_flight)
/// - **退出策略**:监听 `token.cancelled()`,当前正在处理的消息会等待完成
///
/// ## 用法
///
/// ```rust,ignore
/// use sz_rust_core::runtime::queue::{QueueRuntime, QueueRuntimeConfig, QueueConsumer};
/// use sz_orm_queue::{InMemoryQueue, MessageQueue, Message};
/// use async_trait::async_trait;
/// use std::sync::Arc;
/// use tokio_util::sync::CancellationToken;
///
/// struct MyConsumer;
/// #[async_trait]
/// impl QueueConsumer for MyConsumer {
///     async fn handle(&self, msg: &Message) -> Result<(), QueueConsumerError> {
///         println!("got: {:?}", msg.payload);
///         Ok(())
///     }
/// }
///
/// let queue = Arc::new(InMemoryQueue::new(1000));
/// let runtime = QueueRuntime::new(
///     QueueRuntimeConfig::new("orders"),
///     queue,
/// );
/// let token = CancellationToken::new();
/// let handle = runtime.start(Arc::new(MyConsumer), token.clone());
/// // ... 业务运行 ...
/// token.cancel();
/// let _ = handle.await;
/// ```
pub struct QueueRuntime {
    config: QueueRuntimeConfig,
    queue: Arc<dyn MessageQueue>,
}

impl QueueRuntime {
    /// 创建队列运行时
    pub fn new(config: QueueRuntimeConfig, queue: Arc<dyn MessageQueue>) -> Self {
        Self { config, queue }
    }

    /// 启动消费循环(返回 JoinHandle,调用方持有以控制 lifecycle)
    ///
    /// ## 行为
    ///
    /// 1. 每 `poll_interval_ms` 毫秒调用 `queue.consume(topic)` 拉取消息
    /// 2. 收到消息后调用 `consumer.handle(&msg)`
    /// 3. handler 返回 Ok → 自动 `queue.ack(msg.id)`
    /// 4. handler 返回 Err → 跳过 ack(消息留在 in_flight)
    /// 5. 监听 `token.cancelled()`,收到信号后停止拉取新消息
    pub fn start<C>(
        &self,
        consumer: Arc<C>,
        token: CancellationToken,
    ) -> tokio::task::JoinHandle<()>
    where
        C: QueueConsumer + 'static,
    {
        let queue = self.queue.clone();
        let topic = self.config.topic.clone();
        let poll_interval = Duration::from_millis(self.config.poll_interval_ms.max(1));

        tokio::spawn(async move {
            loop {
                tokio::select! {
                    _ = token.cancelled() => break,
                    consume_result = queue.consume(&topic) => {
                        match consume_result {
                            Ok(Some(message)) => {
                                let msg_id = message.id.clone();
                                match consumer.handle(&message).await {
                                    Ok(()) => {
                                        if let Err(e) = queue.ack(&msg_id).await {
                                            tracing::warn!("ack failed for msg {}: {}", msg_id, e);
                                        }
                                    }
                                    Err(e) => {
                                        tracing::warn!(
                                            "consumer handler failed for msg {}: {}",
                                            msg_id,
                                            e
                                        );
                                        // 不 ack,消息留在 in_flight
                                    }
                                }
                            }
                            Ok(None) => {
                                // 队列为空,sleep 后重试
                                tokio::time::sleep(poll_interval).await;
                            }
                            Err(e) => {
                                tracing::error!("queue consume error: {}", e);
                                tokio::time::sleep(poll_interval).await;
                            }
                        }
                    }
                }
            }
        })
    }

    /// 获取主题
    pub fn topic(&self) -> &str {
        &self.config.topic
    }

    /// 获取配置
    pub fn config(&self) -> &QueueRuntimeConfig {
        &self.config
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::orm::{InMemoryQueue, Message, MessageQueue};

    /// 测试用消费者:记录所有处理过的消息 payload
    struct RecordingConsumer {
        payloads: Arc<parking_lot::Mutex<Vec<Vec<u8>>>>,
    }

    impl RecordingConsumer {
        fn new() -> (Self, Arc<parking_lot::Mutex<Vec<Vec<u8>>>>) {
            let payloads = Arc::new(parking_lot::Mutex::new(Vec::new()));
            let consumer = Self {
                payloads: payloads.clone(),
            };
            (consumer, payloads)
        }
    }

    #[async_trait]
    impl QueueConsumer for RecordingConsumer {
        async fn handle(&self, message: &Message) -> Result<(), QueueConsumerError> {
            self.payloads.lock().push(message.payload.clone());
            Ok(())
        }
    }

    /// 失败消费者:总是返回错误
    struct FailingConsumer;

    #[async_trait]
    impl QueueConsumer for FailingConsumer {
        async fn handle(&self, _message: &Message) -> Result<(), QueueConsumerError> {
            Err(QueueConsumerError::Handler("always fail".to_string()))
        }
    }

    /// 创建 InMemoryQueue 并 cast 为 `Arc<dyn MessageQueue>`
    fn make_queue() -> Arc<dyn MessageQueue> {
        Arc::new(InMemoryQueue::new())
    }

    #[test]
    fn test_queue_runtime_config_default() {
        let config = QueueRuntimeConfig::default();
        assert_eq!(config.topic, "default");
        assert_eq!(config.poll_interval_ms, 100);
        assert_eq!(config.max_retries, 0);
    }

    #[test]
    fn test_queue_runtime_config_builder() {
        let config = QueueRuntimeConfig::new("orders")
            .with_poll_interval(50)
            .with_max_retries(3);
        assert_eq!(config.topic, "orders");
        assert_eq!(config.poll_interval_ms, 50);
        assert_eq!(config.max_retries, 3);
    }

    #[test]
    fn test_queue_runtime_topic_accessor() {
        let queue = make_queue();
        let runtime = QueueRuntime::new(QueueRuntimeConfig::new("test"), queue);
        assert_eq!(runtime.topic(), "test");
    }

    #[test]
    fn test_queue_runtime_config_accessor() {
        let queue = make_queue();
        let config = QueueRuntimeConfig::new("test").with_poll_interval(200);
        let runtime = QueueRuntime::new(config, queue);
        assert_eq!(runtime.config().poll_interval_ms, 200);
    }

    #[tokio::test]
    async fn test_consumer_consumes_published_message() {
        let queue = make_queue();
        queue.publish("orders", b"hello").await.unwrap();

        let (consumer, payloads) = RecordingConsumer::new();
        let runtime = QueueRuntime::new(
            QueueRuntimeConfig::new("orders").with_poll_interval(5),
            queue.clone(),
        );

        let token = CancellationToken::new();
        let handle = runtime.start(Arc::new(consumer), token.clone());

        // 等待消费者处理消息
        tokio::time::sleep(Duration::from_millis(100)).await;
        token.cancel();
        let _ = handle.await;

        let recorded = payloads.lock().clone();
        assert_eq!(recorded.len(), 1);
        assert_eq!(recorded[0], b"hello");
    }

    #[tokio::test]
    async fn test_consumer_acks_on_success() {
        let queue = make_queue();
        queue.publish("orders", b"msg1").await.unwrap();

        let (consumer, _payloads) = RecordingConsumer::new();
        let runtime = QueueRuntime::new(
            QueueRuntimeConfig::new("orders").with_poll_interval(5),
            queue.clone(),
        );

        let token = CancellationToken::new();
        let handle = runtime.start(Arc::new(consumer), token.clone());

        tokio::time::sleep(Duration::from_millis(100)).await;
        token.cancel();
        let _ = handle.await;

        // 队列中应该没有 in_flight 消息(已 ack)
        // 注:InMemoryQueue 的 ack 会从 in_flight 移除消息
        // 此处不直接验证 in_flight 状态,因为 InMemoryQueue API 未暴露
        // 通过再次 consume 返回 None 间接验证
        let result = queue.consume("orders").await.unwrap();
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn test_consumer_no_ack_on_failure() {
        let queue = make_queue();
        queue.publish("orders", b"msg1").await.unwrap();

        let runtime = QueueRuntime::new(
            QueueRuntimeConfig::new("orders").with_poll_interval(5),
            queue.clone(),
        );

        let token = CancellationToken::new();
        let handle = runtime.start(Arc::new(FailingConsumer), token.clone());

        tokio::time::sleep(Duration::from_millis(100)).await;
        token.cancel();
        let _ = handle.await;

        // FailingConsumer 不 ack,消息应该留在 in_flight
        // 再次 consume 应该返回 None(因为 in_flight 中的消息不会被重新拉取)
        // 但 in_flight 消息仍占位
        // 注:此行为依赖 InMemoryQueue 的实现
    }

    #[tokio::test]
    async fn test_consumer_stops_on_cancel() {
        let queue = make_queue();
        let (consumer, _payloads) = RecordingConsumer::new();
        let runtime = QueueRuntime::new(
            QueueRuntimeConfig::new("orders").with_poll_interval(5),
            queue,
        );

        let token = CancellationToken::new();
        let handle = runtime.start(Arc::new(consumer), token.clone());

        // 立即 cancel
        token.cancel();
        // 等待任务退出(不应 panic)
        let _ = tokio::time::timeout(Duration::from_millis(500), handle).await;
    }

    #[tokio::test]
    async fn test_consumer_handles_empty_queue() {
        let queue = make_queue();
        let (consumer, payloads) = RecordingConsumer::new();
        let runtime = QueueRuntime::new(
            QueueRuntimeConfig::new("empty").with_poll_interval(5),
            queue.clone(),
        );

        let token = CancellationToken::new();
        let handle = runtime.start(Arc::new(consumer), token.clone());

        // 等待一段时间,队列始终为空
        tokio::time::sleep(Duration::from_millis(50)).await;
        token.cancel();
        let _ = handle.await;

        // 没有消息被处理
        assert!(payloads.lock().is_empty());
    }

    #[tokio::test]
    async fn test_consumer_processes_multiple_messages() {
        let queue = make_queue();
        // 发布 3 条消息
        queue.publish("orders", b"msg1").await.unwrap();
        queue.publish("orders", b"msg2").await.unwrap();
        queue.publish("orders", b"msg3").await.unwrap();

        let (consumer, payloads) = RecordingConsumer::new();
        let runtime = QueueRuntime::new(
            QueueRuntimeConfig::new("orders").with_poll_interval(5),
            queue,
        );

        let token = CancellationToken::new();
        let handle = runtime.start(Arc::new(consumer), token.clone());

        // 等待所有消息被处理
        tokio::time::sleep(Duration::from_millis(200)).await;
        token.cancel();
        let _ = handle.await;

        let recorded = payloads.lock().clone();
        assert_eq!(recorded.len(), 3);
    }

    #[test]
    fn test_queue_consumer_error_variants() {
        let handler_err = QueueConsumerError::Handler("test".to_string());
        let queue_err = QueueConsumerError::Queue("queue fail".to_string());
        assert!(format!("{}", handler_err).contains("consumer error"));
        assert!(format!("{}", queue_err).contains("queue error"));
    }
}