sz_rust_core/
runtime.rs

1//! SZ-Rust Runtime — Swoole/Worker 适配主入口
2//!
3//! ## PHP 对齐
4//!
5//! 对齐 PHP `topthink/think-swoole` 的运行时模型：
6//!
7//! | PHP Swoole | Rust SZ-Rust | 说明 |
8//! |------------|--------------|------|
9//! | `Swoole\Runtime::enableCoroutine()` | `SzRuntime::block_on(fut)` | 协程入口 |
10//! | `swoole_cpu_num()` | `num_cpus::get()` | 默认 worker 数 |
11//! | `worker_num` 配置项 | `SzRuntime::with_worker_threads(n)` | 自定义 worker 数 |
12//! | `Swoole\Process::signal()` | `tokio::signal` + `CancellationToken` | 信号处理 |
13//! | `Swoole\Timer::tick()` | `tokio::time::interval` + `tokio::select!` | 定时任务 |
14//! | `swoole_event::defer()` | `SzRuntime::spawn(fut)` | 异步任务 |
15//!
16//! ## 模块结构
17//!
18//! | 模块 | 对齐 PHP | 子任务 |
19//! |------|---------|--------|
20//! | `runtime::worker` | `worker_num` 配置 | 9.2 |
21//! | `runtime::spawn` | `swoole_event::defer` | 9.3 |
22//! | `runtime::queue` | `think-queue` 消费者 | 9.4 |
23//! | `runtime::mqtt` | 长连接 | 9.5 |
24//! | `runtime::websocket` | `think-worker` | 9.6 |
25//! | `runtime::scheduler` | `Crontab` | 9.7 |
26//! | `runtime::shutdown` | 优雅关闭 | 9.8 |
27//! | `runtime::signal` | `SIGTERM/SIGINT` | 9.9 |
28//!
29//! ## 关键决策
30//!
31//! - **不修改 `sz-orm-scheduler` 内部 API**：保持 `CronScheduler::start()`/`stop()` 向后兼容，
32//!   在 sz-rust 侧用 `tokio::time::interval` + `try_fire_due()` 重写循环。
33//! - **使用 `CancellationToken` 统一关闭广播**：替代 `AtomicBool` + `oneshot`，支持父子层级。
34//! - **双平台信号处理**：Unix 用 `SignalKind::terminate/interrupt`，Windows 用 `ctrl_c/ctrl_close`。
35
36pub mod mqtt;
37pub mod queue;
38pub mod scheduler;
39pub mod shutdown;
40pub mod signal;
41pub mod spawn;
42pub mod websocket;
43pub mod worker;
44
45pub use mqtt::{MqttRuntime, MqttRuntimeConfig};
46pub use queue::{QueueConsumer, QueueRuntime, QueueRuntimeConfig};
47pub use scheduler::SchedulerRuntime;
48pub use shutdown::GracefulShutdown;
49pub use signal::shutdown_signal;
50pub use spawn::spawn_with_token;
51pub use websocket::{WebSocketRuntime, WebSocketRuntimeConfig};
52pub use worker::WorkerConfig;
53
54use std::future::Future;
55use std::time::Duration;
56
57use tokio_util::sync::CancellationToken;
58
59/// SZ-Rust 异步运行时
60///
61/// 对齐 PHP `think-swoole` 的运行时模型：基于 tokio multi_thread runtime，
62/// worker 数量默认 = CPU 核数（对齐 `swoole_cpu_num()`）。
63///
64/// ## 设计
65///
66/// - 封装 `tokio::runtime::Runtime`，对外暴露 `block_on` / `spawn` 入口
67/// - 持有 `CancellationToken`，所有后台任务通过 `shutdown_token()` 监听关闭信号
68/// - `shutdown_timeout` 触发关闭流程：先 `cancel()` 通知所有任务，等待超时后 drop runtime
69///
70/// ## 用法
71///
72/// ```rust,ignore
73/// use sz_rust_core::runtime::SzRuntime;
74/// use std::time::Duration;
75///
76/// let rt = SzRuntime::new();
77/// assert!(rt.worker_threads() > 0);
78///
79/// // spawn 后台任务
80/// let handle = rt.spawn(async { 42 });
81/// assert_eq!(rt.block_on(handle).unwrap(), 42);
82///
83/// // 优雅关闭
84/// assert!(rt.shutdown_timeout(Duration::from_millis(100)));
85/// ```
86pub struct SzRuntime {
87    /// 内部 tokio runtime（multi_thread）
88    runtime: tokio::runtime::Runtime,
89    /// worker 线程数（对齐 swoole `worker_num`）
90    worker_threads: usize,
91    /// 关闭令牌：所有后台任务通过 `shutdown_token()` 获取子 token 监听关闭
92    shutdown_token: CancellationToken,
93}
94
95impl SzRuntime {
96    /// 创建 SzRuntime，worker_threads = `num_cpus::get()`（对齐 `swoole_cpu_num()`）
97    pub fn new() -> Self {
98        Self::with_worker_threads(num_cpus::get())
99    }
100
101    /// 创建 SzRuntime，自定义 worker_threads（对齐 `worker_num` 配置项）
102    ///
103    /// - `worker_threads = 0` 会被强制为 1
104    pub fn with_worker_threads(worker_threads: usize) -> Self {
105        let n = worker_threads.max(1);
106        let runtime = tokio::runtime::Builder::new_multi_thread()
107            .worker_threads(n)
108            .enable_all()
109            .thread_name("sz-rust-worker")
110            .build()
111            .expect("Failed to create tokio runtime");
112        Self {
113            runtime,
114            worker_threads: n,
115            shutdown_token: CancellationToken::new(),
116        }
117    }
118
119    /// 获取 worker 线程数
120    pub fn worker_threads(&self) -> usize {
121        self.worker_threads
122    }
123
124    /// 获取关闭令牌的克隆
125    ///
126    /// 后台任务通过 `rt.shutdown_token().cancelled().await` 监听关闭信号。
127    pub fn shutdown_token(&self) -> CancellationToken {
128        self.shutdown_token.clone()
129    }
130
131    /// 在 runtime 上 spawn 异步任务（对齐 `swoole_event::defer`）
132    ///
133    /// 返回 `JoinHandle`，可 await 获取结果。
134    pub fn spawn<F>(&self, future: F) -> tokio::task::JoinHandle<F::Output>
135    where
136        F: Future + Send + 'static,
137        F::Output: Send + 'static,
138    {
139        self.runtime.spawn(future)
140    }
141
142    /// 在 runtime 上阻塞运行 future（对齐 `Swoole\Runtime::enableCoroutine` 入口）
143    ///
144    /// 调用线程会阻塞直到 future 完成。
145    pub fn block_on<F>(&self, future: F) -> F::Output
146    where
147        F: Future,
148    {
149        self.runtime.block_on(future)
150    }
151
152    /// 触发优雅关闭：`cancel()` 通知所有后台任务，等待 `timeout` 后 drop runtime
153    ///
154    /// - 返回 `true`：runtime 已关闭
155    /// - 注意：Drop runtime 时 tokio 会等待所有任务完成，超时由调用方控制
156    pub fn shutdown_timeout(self, timeout: Duration) -> bool {
157        self.shutdown_token.cancel();
158        // 给后台任务响应关闭信号的时间
159        self.runtime.block_on(async {
160            let _ = tokio::time::timeout(timeout, async {
161                // 等待一小段时间让任务响应 cancel
162                tokio::time::sleep(Duration::from_millis(10)).await;
163            })
164            .await;
165        });
166        // Drop runtime 会等待所有任务完成（可能阻塞）
167        drop(self.runtime);
168        true
169    }
170
171    /// 获取内部 runtime 句柄（用于在不持有 SzRuntime 时 spawn 任务）
172    pub fn handle(&self) -> tokio::runtime::Handle {
173        self.runtime.handle().clone()
174    }
175}
176
177impl Default for SzRuntime {
178    fn default() -> Self {
179        Self::new()
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186
187    #[test]
188    fn test_new_default_worker_threads() {
189        let rt = SzRuntime::new();
190        assert_eq!(rt.worker_threads(), num_cpus::get());
191    }
192
193    #[test]
194    fn test_with_worker_threads_custom() {
195        let rt = SzRuntime::with_worker_threads(2);
196        assert_eq!(rt.worker_threads(), 2);
197    }
198
199    #[test]
200    fn test_with_worker_threads_zero_falls_back_to_one() {
201        let rt = SzRuntime::with_worker_threads(0);
202        assert_eq!(rt.worker_threads(), 1);
203    }
204
205    #[test]
206    fn test_spawn_and_block_on() {
207        let rt = SzRuntime::with_worker_threads(1);
208        let handle = rt.spawn(async { 42 });
209        let result = rt.block_on(handle).unwrap();
210        assert_eq!(result, 42);
211    }
212
213    #[test]
214    fn test_block_on_directly() {
215        let rt = SzRuntime::with_worker_threads(1);
216        let result = rt.block_on(async { 100 });
217        assert_eq!(result, 100);
218    }
219
220    #[test]
221    fn test_shutdown_token_cancellation() {
222        let rt = SzRuntime::with_worker_threads(1);
223        let token = rt.shutdown_token();
224        assert!(!token.is_cancelled());
225        assert!(rt.shutdown_timeout(Duration::from_millis(50)));
226        // shutdown_timeout 已消费 rt，token 已 cancel
227        assert!(token.is_cancelled());
228    }
229
230    #[test]
231    fn test_spawn_with_token_cancellation() {
232        let rt = SzRuntime::with_worker_threads(1);
233        let token = rt.shutdown_token();
234        let handle = rt.spawn(async move {
235            // 模拟后台任务：监听 cancel
236            token.cancelled().await;
237            99
238        });
239        // 触发关闭
240        let token2 = rt.shutdown_token();
241        token2.cancel();
242        let result = rt.block_on(handle).unwrap();
243        assert_eq!(result, 99);
244    }
245
246    #[test]
247    fn test_handle_can_spawn() {
248        let rt = SzRuntime::with_worker_threads(1);
249        let handle = rt.handle();
250        let task = handle.spawn(async { 7 });
251        let result = rt.block_on(task).unwrap();
252        assert_eq!(result, 7);
253    }
254
255    #[test]
256    fn test_default_impl_equals_new() {
257        let rt1 = SzRuntime::default();
258        let rt2 = SzRuntime::new();
259        assert_eq!(rt1.worker_threads(), rt2.worker_threads());
260    }
261
262    #[test]
263    fn test_multiple_runtime_instances() {
264        // 验证可以创建多个独立的 runtime 实例
265        let rt1 = SzRuntime::with_worker_threads(1);
266        let rt2 = SzRuntime::with_worker_threads(1);
267        let h1 = rt1.spawn(async { 1 });
268        let h2 = rt2.spawn(async { 2 });
269        assert_eq!(rt1.block_on(h1).unwrap(), 1);
270        assert_eq!(rt2.block_on(h2).unwrap(), 2);
271    }
272}
sz_rust_core/runtime.rs

sz_rust_core/
runtime.rs
