1use std::{
69 collections::{BTreeMap, HashMap},
70 ops::{AddAssign, Sub},
71 time::{Duration, Instant},
72};
73
74#[inline(always)]
75fn tv_to_duration(tv: libc::timeval) -> Duration {
76 Duration::new(tv.tv_sec as u64, (tv.tv_usec * 1000) as u32)
77}
78
79#[derive(Clone, Copy)]
80struct ResourceUsage(libc::rusage);
81
82impl ResourceUsage {
83 pub fn now() -> Self {
84 unsafe {
85 let mut rusage = std::mem::MaybeUninit::<libc::rusage>::uninit();
86 assert!(0 == libc::getrusage(libc::RUSAGE_SELF, rusage.as_mut_ptr()));
87 Self(rusage.assume_init())
88 }
89 }
90
91 #[inline(always)]
92 pub fn utime(&self) -> Duration {
93 tv_to_duration(self.0.ru_utime)
94 }
95
96 #[inline(always)]
97 pub fn stime(&self) -> Duration {
98 tv_to_duration(self.0.ru_stime)
99 }
100
101 #[inline(always)]
102 pub fn cpu_time(&self) -> Duration {
103 self.utime() + self.stime()
104 }
105}
106
107impl std::fmt::Debug for ResourceUsage {
108 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109 f.debug_struct("ResourceUsage")
110 .field("ru_utime.tv_sec", &self.0.ru_utime.tv_sec)
111 .field("ru_utime.tv_usec", &self.0.ru_utime.tv_usec)
112 .field("ru_stime.tv_sec", &self.0.ru_stime.tv_sec)
113 .field("ru_stime.tv_usec", &self.0.ru_stime.tv_usec)
114 .finish()
115 }
116}
117
118#[derive(Debug, Clone, Copy)]
119struct UsageSnapshot {
120 pub time: Instant,
121 pub usage: ResourceUsage,
122}
123
124impl UsageSnapshot {
125 pub fn now() -> Self {
126 Self {
127 time: Instant::now(),
128 usage: ResourceUsage::now(),
129 }
130 }
131}
132
133impl Sub for UsageSnapshot {
134 type Output = UsageMeasurement;
135
136 fn sub(self, rhs: Self) -> Self::Output {
137 UsageMeasurement {
138 utime: self.usage.utime() - rhs.usage.utime(),
139 stime: self.usage.stime() - rhs.usage.stime(),
140 wtime: self.time.duration_since(rhs.time),
141 }
142 }
143}
144
145#[derive(Debug)]
146struct ScopeFrame {
147 path: String,
148 start: UsageSnapshot,
149}
150
151#[derive(Debug, Default, Clone, Copy)]
152pub struct UsageMeasurement {
153 utime: Duration,
155 stime: Duration,
158 wtime: Duration,
160}
161
162impl UsageMeasurement {}
163
164impl AddAssign for UsageMeasurement {
165 fn add_assign(&mut self, rhs: Self) {
166 self.utime += rhs.utime;
167 self.stime += rhs.stime;
168 self.wtime += rhs.wtime;
169 }
170}
171
172#[derive(Debug)]
173pub struct ResourceMeterStack {
174 store: HashMap<String, UsageMeasurement>,
175 stack: Vec<ScopeFrame>,
176 last_snapshot: UsageSnapshot,
177}
178
179impl ResourceMeterStack {
180 pub fn new() -> Self {
181 Self {
182 store: HashMap::new(),
183 stack: vec![ScopeFrame {
184 path: String::new(),
185 start: UsageSnapshot::now(),
186 }],
187 last_snapshot: UsageSnapshot::now(),
188 }
189 }
190
191 pub fn push<T: Into<String>>(&mut self, key: T) {
192 self.stack.push(ScopeFrame {
193 path: key.into(),
194 start: UsageSnapshot::now(),
195 });
196 }
197
198 pub fn pop(&mut self) {
199 let now = UsageSnapshot::now();
200 let curr = self.stack.pop().unwrap();
201 let slot = self.store.entry(curr.path).or_default();
202 *slot += now - curr.start;
203 self.last_snapshot = now;
204 }
205
206 pub fn into_report(mut self, format: ReportFormat) -> Report {
207 while !self.stack.is_empty() {
208 self.pop();
209 }
210 match format {
211 ReportFormat::Flat => Report::Flat(self.store),
212 ReportFormat::Tree => {
213 let mut root = TreeNode::default();
214 for (path, usage) in self.store {
215 root.insert(path, usage);
216 }
217 Report::Tree(root)
218 }
219 }
220 }
221
222 pub fn finish(mut self) -> TreeReport {
223 while !self.stack.is_empty() {
224 self.pop();
225 }
226 TreeReport::new(self.store)
227 }
228}
229
230pub enum ReportFormat {
231 Flat,
232 Tree,
233}
234
235pub enum Report {
236 Flat(HashMap<String, UsageMeasurement>),
237 Tree(TreeNode),
238}
239
240#[derive(Debug)]
241pub struct TreeReport(TreeNode);
242
243impl TreeReport {
244 pub fn new(measurements: HashMap<String, UsageMeasurement>) -> Self {
245 let mut root = TreeNode::default();
246 for (path, usage) in measurements {
247 root.insert(path, usage);
248 }
249 Self(root)
250 }
251}
252
253#[derive(Debug, Default)]
254struct TreeNode {
255 usage: Option<UsageMeasurement>,
256 children: BTreeMap<String, TreeNode>,
257}
258
259impl TreeNode {
260 fn insert(&mut self, path: String, usage: UsageMeasurement) {
261 let mut current = self;
262 for part in path.split('/') {
263 current = current.children.entry(part.to_string()).or_default();
264 }
265 current.usage = Some(usage);
266 }
267
268 fn fmt_with_indent(&self, f: &mut std::fmt::Formatter<'_>, depth: usize) -> std::fmt::Result {
269 let width = 20;
270 for (name, child) in &self.children {
271 let indent = " ".repeat(depth);
272
273 if let Some(usage) = child.usage {
274 writeln!(
275 f,
276 "{indent}{:<} real: {:>6?}, user: {:>6?}, sys: {:?}",
277 name,
278 usage.wtime.as_millis(),
279 usage.utime.as_millis(),
280 usage.stime.as_millis(),
281 )?;
282 } else {
283 writeln!(f, "{indent}{name}")?;
284 }
285 child.fmt_with_indent(f, depth + 1)?;
286 }
287 Ok(())
288 }
289}
290
291impl std::fmt::Display for TreeReport {
295 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
296 writeln!(f, "--- Resource Usage Report ---")?;
297 self.0.fmt_with_indent(f, 0)?;
298 Ok(())
299 }
300}
301
302#[cfg(test)]
303mod tests {
304 use super::*;
305
306 #[test]
307 fn test_simple() {
308 let mut stk = ResourceMeterStack::new();
309 let now = Instant::now();
310 std::thread::sleep(Duration::from_millis(200));
311 stk.push("sleep600");
312 stk.push("sleep600/300");
313 std::thread::sleep(Duration::from_millis(300));
314 stk.pop();
315 std::thread::sleep(Duration::from_millis(100));
316 stk.push("sleep600/200");
317 std::thread::sleep(Duration::from_millis(200));
318 stk.pop();
319 stk.pop();
320 std::thread::sleep(Duration::from_millis(200));
321 let report = stk.finish();
322 eprintln!("{}, took: {:?}", report, now.elapsed());
323 }
324
325 #[test]
326 fn test_resource_usage() {
327 let rusage = ResourceUsage::now();
328 eprintln!("{:#?}", rusage);
329 }
330}