aster/sandbox/
resource_limits.rs

1//! 资源限制
2//!
3//! 提供进程资源限制和使用监控
4
5use serde::{Deserialize, Serialize};
6use std::time::{Duration, Instant};
7
8/// 资源使用情况
9#[derive(Debug, Clone, Default, Serialize, Deserialize)]
10pub struct ResourceUsage {
11    /// 内存使用（字节）
12    pub memory_bytes: u64,
13    /// CPU 使用率 (0-100)
14    pub cpu_percent: f32,
15    /// 进程数
16    pub process_count: u32,
17    /// 打开的文件描述符数
18    pub file_descriptors: u32,
19    /// 执行时间（毫秒）
20    pub execution_time_ms: u64,
21}
22
23/// 资源限制器
24pub struct ResourceLimiter {
25    /// 最大内存
26    max_memory: Option<u64>,
27    /// 最大 CPU
28    max_cpu: Option<u32>,
29    /// 最大进程数
30    max_processes: Option<u32>,
31    /// 最大执行时间
32    max_execution_time: Option<Duration>,
33    /// 最大文件描述符
34    max_file_descriptors: Option<u32>,
35    /// 开始时间
36    start_time: Option<Instant>,
37}
38
39impl ResourceLimiter {
40    /// 创建新的资源限制器
41    pub fn new() -> Self {
42        Self {
43            max_memory: None,
44            max_cpu: None,
45            max_processes: None,
46            max_execution_time: None,
47            max_file_descriptors: None,
48            start_time: None,
49        }
50    }
51
52    /// 从配置创建
53    pub fn from_limits(limits: &super::config::ResourceLimits) -> Self {
54        Self {
55            max_memory: limits.max_memory,
56            max_cpu: limits.max_cpu,
57            max_processes: limits.max_processes,
58            max_execution_time: limits.max_execution_time.map(Duration::from_millis),
59            max_file_descriptors: limits.max_file_descriptors,
60            start_time: None,
61        }
62    }
63
64    /// 设置最大内存
65    pub fn with_max_memory(mut self, bytes: u64) -> Self {
66        self.max_memory = Some(bytes);
67        self
68    }
69
70    /// 设置最大 CPU
71    pub fn with_max_cpu(mut self, percent: u32) -> Self {
72        self.max_cpu = Some(percent);
73        self
74    }
75
76    /// 设置最大进程数
77    pub fn with_max_processes(mut self, count: u32) -> Self {
78        self.max_processes = Some(count);
79        self
80    }
81
82    /// 设置最大执行时间
83    pub fn with_max_execution_time(mut self, duration: Duration) -> Self {
84        self.max_execution_time = Some(duration);
85        self
86    }
87
88    /// 开始计时
89    pub fn start(&mut self) {
90        self.start_time = Some(Instant::now());
91    }
92
93    /// 检查是否超时
94    pub fn is_timeout(&self) -> bool {
95        if let (Some(start), Some(max_time)) = (self.start_time, self.max_execution_time) {
96            return start.elapsed() > max_time;
97        }
98        false
99    }
100
101    /// 检查资源使用是否超限
102    pub fn check_limits(&self, usage: &ResourceUsage) -> Result<(), ResourceLimitError> {
103        if let Some(max_memory) = self.max_memory {
104            if usage.memory_bytes > max_memory {
105                return Err(ResourceLimitError::MemoryExceeded {
106                    used: usage.memory_bytes,
107                    limit: max_memory,
108                });
109            }
110        }
111
112        if let Some(max_cpu) = self.max_cpu {
113            if usage.cpu_percent > max_cpu as f32 {
114                return Err(ResourceLimitError::CpuExceeded {
115                    used: usage.cpu_percent,
116                    limit: max_cpu as f32,
117                });
118            }
119        }
120
121        if let Some(max_processes) = self.max_processes {
122            if usage.process_count > max_processes {
123                return Err(ResourceLimitError::ProcessesExceeded {
124                    used: usage.process_count,
125                    limit: max_processes,
126                });
127            }
128        }
129
130        if let Some(max_fds) = self.max_file_descriptors {
131            if usage.file_descriptors > max_fds {
132                return Err(ResourceLimitError::FileDescriptorsExceeded {
133                    used: usage.file_descriptors,
134                    limit: max_fds,
135                });
136            }
137        }
138
139        if self.is_timeout() {
140            return Err(ResourceLimitError::Timeout {
141                elapsed: self.start_time.map(|s| s.elapsed()).unwrap_or_default(),
142                limit: self.max_execution_time.unwrap_or_default(),
143            });
144        }
145
146        Ok(())
147    }
148
149    /// 获取剩余执行时间
150    pub fn remaining_time(&self) -> Option<Duration> {
151        match (self.start_time, self.max_execution_time) {
152            (Some(start), Some(max_time)) => {
153                let elapsed = start.elapsed();
154                if elapsed < max_time {
155                    Some(max_time - elapsed)
156                } else {
157                    Some(Duration::ZERO)
158                }
159            }
160            _ => None,
161        }
162    }
163}
164
165impl Default for ResourceLimiter {
166    fn default() -> Self {
167        Self::new()
168    }
169}
170
171/// 资源限制错误
172#[derive(Debug, Clone)]
173pub enum ResourceLimitError {
174    /// 内存超限
175    MemoryExceeded { used: u64, limit: u64 },
176    /// CPU 超限
177    CpuExceeded { used: f32, limit: f32 },
178    /// 进程数超限
179    ProcessesExceeded { used: u32, limit: u32 },
180    /// 文件描述符超限
181    FileDescriptorsExceeded { used: u32, limit: u32 },
182    /// 执行超时
183    Timeout { elapsed: Duration, limit: Duration },
184}
185
186impl std::fmt::Display for ResourceLimitError {
187    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
188        match self {
189            Self::MemoryExceeded { used, limit } => {
190                write!(f, "内存超限: 使用 {} 字节，限制 {} 字节", used, limit)
191            }
192            Self::CpuExceeded { used, limit } => {
193                write!(f, "CPU 超限: 使用 {:.1}%，限制 {:.1}%", used, limit)
194            }
195            Self::ProcessesExceeded { used, limit } => {
196                write!(f, "进程数超限: 使用 {}，限制 {}", used, limit)
197            }
198            Self::FileDescriptorsExceeded { used, limit } => {
199                write!(f, "文件描述符超限: 使用 {}，限制 {}", used, limit)
200            }
201            Self::Timeout { elapsed, limit } => {
202                write!(f, "执行超时: 已执行 {:?}，限制 {:?}", elapsed, limit)
203            }
204        }
205    }
206}
207
208impl std::error::Error for ResourceLimitError {}
209
210/// 构建 ulimit 参数
211pub fn build_ulimit_args(limits: &super::config::ResourceLimits) -> Vec<String> {
212    let mut args = Vec::new();
213
214    if let Some(max_memory) = limits.max_memory {
215        // 虚拟内存限制 (KB)
216        args.push(format!("-v {}", max_memory / 1024));
217    }
218
219    if let Some(max_fds) = limits.max_file_descriptors {
220        // 文件描述符限制
221        args.push(format!("-n {}", max_fds));
222    }
223
224    if let Some(max_processes) = limits.max_processes {
225        // 进程数限制
226        args.push(format!("-u {}", max_processes));
227    }
228
229    if let Some(max_file_size) = limits.max_file_size {
230        // 文件大小限制 (KB)
231        args.push(format!("-f {}", max_file_size / 1024));
232    }
233
234    args
235}
aster/sandbox/resource_limits.rs

aster/sandbox/
resource_limits.rs
