sz_rust_cli/cmd/serve/
runtime.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2024-2026 SZ-Rust Team
3//
4
5//! 多 worker runtime 构建
6//!
7//! 根据 worker 数量构建 tokio multi-thread runtime，用于 serve 命令启动 HTTP 服务。
8
9use crate::error::CliError;
10
11/// 校验 worker 数量合法性
12///
13/// - `workers == 0` → 错误
14/// - `workers > 1024` → 错误
15pub fn validate_workers(workers: u16) -> Result<(), CliError> {
16    if workers == 0 {
17        return Err(CliError::Generic("worker 数量必须 >= 1".to_string()));
18    }
19    if workers > 1024 {
20        return Err(CliError::Generic("worker 数量超过上限 1024".to_string()));
21    }
22    Ok(())
23}
24
25/// 构建 multi-thread tokio runtime
26///
27/// 指定 worker 线程数，启用所有功能（IO + time + 等）。
28pub fn build_runtime(workers: u16) -> Result<tokio::runtime::Runtime, CliError> {
29    validate_workers(workers)?;
30    tokio::runtime::Builder::new_multi_thread()
31        .worker_threads(workers as usize)
32        .enable_all()
33        .build()
34        .map_err(|e| CliError::Generic(format!("tokio runtime 构建失败: {e}")))
35}
36
37/// 解析有效 worker 数
38///
39/// 优先级：CLI 参数 > 配置文件 > CPU 核心数
40pub fn resolve_workers(cli_workers: Option<u16>, config_workers: u16) -> u16 {
41    let w = cli_workers.unwrap_or(config_workers);
42    if w == 0 {
43        std::thread::available_parallelism()
44            .map(|n| n.get() as u16)
45            .unwrap_or(1)
46    } else {
47        w
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    #[test]
56    fn test_validate_workers_zero_rejected() {
57        assert!(validate_workers(0).is_err());
58    }
59
60    #[test]
61    fn test_validate_workers_one_ok() {
62        assert!(validate_workers(1).is_ok());
63    }
64
65    #[test]
66    fn test_validate_workers_over_limit_rejected() {
67        assert!(validate_workers(1025).is_err());
68    }
69
70    #[test]
71    fn test_validate_workers_at_limit_ok() {
72        assert!(validate_workers(1024).is_ok());
73    }
74
75    #[test]
76    fn test_build_runtime_one_worker() {
77        assert!(build_runtime(1).is_ok());
78    }
79
80    #[test]
81    fn test_build_runtime_zero_rejected() {
82        assert!(build_runtime(0).is_err());
83    }
84
85    #[test]
86    fn test_build_runtime_two_workers() {
87        assert!(build_runtime(2).is_ok());
88    }
89
90    #[test]
91    fn test_resolve_workers_cli_priority() {
92        assert_eq!(resolve_workers(Some(4), 8), 4);
93    }
94
95    #[test]
96    fn test_resolve_workers_config_fallback() {
97        assert_eq!(resolve_workers(None, 8), 8);
98    }
99
100    #[test]
101    fn test_resolve_workers_zero_falls_to_cpu() {
102        let resolved = resolve_workers(None, 0);
103        assert!(resolved >= 1, "CPU 核心数至少为 1");
104    }
105}
sz_rust_cli/cmd/serve/runtime.rs

sz_rust_cli/cmd/serve/
runtime.rs
