1use std::{
2 cmp::min,
3 time::{Duration, Instant, SystemTime, UNIX_EPOCH},
4};
5
6use rust_decimal::Decimal;
7
8use super::Average;
9
10#[derive(Debug, Clone, Copy)]
11pub enum StatPeriod {
12 Live(Instant), Interval(Instant, Instant), }
15
16impl Default for StatPeriod {
17 fn default() -> Self {
18 StatPeriod::Live(Instant::now())
19 }
20}
21
22impl StatPeriod {
23 pub fn period(&self) -> Duration {
24 match self {
25 StatPeriod::Live(start) => Instant::now().duration_since(*start),
26 StatPeriod::Interval(start, end) => end.duration_since(*start),
27 }
28 }
29
30 pub fn is_global(&self) -> bool {
31 matches!(self, StatPeriod::Live(_))
32 }
33
34 pub fn is_period(&self) -> bool {
35 matches!(self, StatPeriod::Interval(_, _))
36 }
37
38 pub fn start(&self) -> Instant {
39 match self {
40 StatPeriod::Live(start) => *start,
41 StatPeriod::Interval(start, _) => *start,
42 }
43 }
44
45 pub fn end(&self) -> Instant {
46 match self {
47 StatPeriod::Live(_) => Instant::now(),
48 StatPeriod::Interval(_, end) => *end,
49 }
50 }
51
52 fn format_duration(from: &Instant, to: &Instant) -> String {
53 let duration = to.duration_since(*from);
54 let secs = duration.as_secs();
55 let hours = secs / 3600;
56 let minutes = (secs % 3600) / 60;
57 let seconds = secs % 60;
58
59 if hours > 0 {
60 format!("{}h {}m {}s", hours, minutes, seconds)
61 } else if minutes > 0 {
62 format!("{}m {}s", minutes, seconds)
63 } else {
64 format!("{}s", seconds)
65 }
66 }
67
68 fn format_instant(instant: &Instant) -> String {
69 let system_now = SystemTime::now();
71 let instant_now = Instant::now();
72
73 let duration_since_instant = instant_now.duration_since(*instant);
74 let system_time = system_now - duration_since_instant;
75
76 let duration_since_epoch = system_time
78 .duration_since(UNIX_EPOCH)
79 .expect("Time went backwards");
80
81 let total_secs = duration_since_epoch.as_secs();
82
83 let seconds_in_day = total_secs % 86400;
85 let hours = seconds_in_day / 3600;
86 let minutes = (seconds_in_day % 3600) / 60;
87 let seconds = seconds_in_day % 60;
88
89 format!("{:02}:{:02}:{:02}", hours, minutes, seconds)
90 }
91
92 pub fn format(&self) -> String {
93 match self {
94 StatPeriod::Live(start) => {
95 format!(
96 "[last {}]",
97 StatPeriod::format_duration(start, &Instant::now())
98 )
99 }
100 StatPeriod::Interval(start, end) => {
101 format!(
102 "[{} +{}]",
103 StatPeriod::format_instant(start),
104 StatPeriod::format_duration(start, end)
105 )
106 }
107 }
108 }
109
110 pub fn format_rps(&self, count: usize) -> String {
111 let duration = self.period();
112 let secs = duration.as_secs_f64();
113 if secs > 0.0 {
114 let rps = count as f64 / secs;
115 format!("{:.2} rps", rps)
116 } else {
117 "N/A".to_string()
118 }
119 }
120}
121
122#[derive(Debug, Clone, Default, Copy)]
123pub struct Stats {
124 pub period: StatPeriod,
125
126 pub success: usize,
127 pub retries: usize,
128 pub errors: usize,
129
130 pub average_latency: Average,
131}
132
133impl Stats {
134 pub fn new() -> Self {
135 Self::default()
136 }
137
138 pub fn empty_interval() -> Self {
139 let now = Instant::now();
140 Self {
141 period: StatPeriod::Interval(now, now),
142 ..Default::default()
143 }
144 }
145
146 pub fn sleep_secs(&mut self, secs: u64) -> Duration {
147 let target_time = min(self.period.end(), Instant::now()) + Duration::from_secs(secs);
149 if target_time < Instant::now() {
150 Duration::ZERO
151 } else {
152 target_time.duration_since(Instant::now())
153 }
154 }
155
156 pub fn snapshot(&self) -> Stats {
157 if self.period.is_period() {
158 return *self;
159 }
160
161 Stats {
162 period: StatPeriod::Interval(self.period.start(), Instant::now()),
163 success: self.success,
164 retries: self.retries,
165 errors: self.errors,
166 average_latency: self.average_latency,
167 }
168 }
169
170 pub fn success(&mut self, latency: Decimal) {
172 self.success += 1;
173 self.average_latency.add_sample(latency);
174 }
175
176 pub fn retry(&mut self) {
178 self.retries += 1;
179 }
180
181 pub fn error(&mut self) {
183 self.errors += 1;
184 }
185
186 pub fn get_total_requests(&self) -> usize {
187 self.success + self.retries + self.errors
188 }
189
190 pub fn get_success(&self) -> usize {
191 self.success
192 }
193
194 pub fn get_retries(&self) -> usize {
195 self.retries
196 }
197
198 pub fn get_errors(&self) -> usize {
199 self.errors
200 }
201
202 pub fn format(&self) -> String {
203 let mut str = self.period.format();
204
205 str.push_str(&format!(
206 " {} requests ({})",
207 self.success,
208 self.period.format_rps(self.success)
209 ));
210
211 if self.success > 0 {
212 str.push_str(&format!(
213 ", avg latency: {:.2} ms",
214 self.average_latency.get_value()
215 ));
216 }
217
218 if self.retries > 0 {
219 str.push_str(&format!(", {} retries", self.retries));
220 }
221
222 if self.errors > 0 {
223 str.push_str(&format!(", {} errors", self.errors));
224 }
225
226 str
227 }
228
229 pub fn format_notime(&self) -> String {
230 let mut str = String::new();
231
232 str.push_str(&format!("{} requests", self.success,));
233
234 if self.success > 0 {
235 str.push_str(&format!(
236 ", avg latency: {:.2} ms",
237 self.average_latency.get_value()
238 ));
239 }
240 if self.retries > 0 {
241 str.push_str(&format!(", {} retries", self.retries));
242 }
243
244 if self.errors > 0 {
245 str.push_str(&format!(", {} errors", self.errors));
246 }
247
248 str
249 }
250}
251
252impl std::ops::Add for Stats {
253 type Output = Stats;
254
255 fn add(self, other: Stats) -> Stats {
256 let period = StatPeriod::Interval(
257 self.period.start(),
258 self.period.end() + other.period.end().duration_since(other.period.start()),
259 );
260
261 Stats {
264 period,
265
266 success: self.success + other.success,
267 retries: self.retries + other.retries,
268 errors: self.errors + other.errors,
269
270 average_latency: self.average_latency + other.average_latency,
271 }
272 }
273}
274
275impl std::ops::Sub for Stats {
276 type Output = Stats;
277
278 fn sub(self, other: Stats) -> Stats {
279 let period = StatPeriod::Interval(
280 self.period.start() + other.period.end().duration_since(other.period.start()),
281 self.period.end(),
282 );
283
284 Stats {
285 period,
286
287 success: self.success.saturating_sub(other.success),
288 retries: self.retries.saturating_sub(other.retries),
289 errors: self.errors.saturating_sub(other.errors),
290
291 average_latency: self.average_latency - other.average_latency,
292 }
293 }
294}
295
296#[cfg(test)]
297mod tests {
298 use rust_decimal::dec;
299
300 use super::*;
301 #[test]
302 fn test_stat_period() {
303 let mut stats = Stats::new();
305
306 stats.success(dec!(0.1));
308 stats.success(dec!(0.2));
309
310 stats.retry();
312
313 let minute1 = stats.snapshot();
315
316 stats.success(dec!(0.3));
318 stats.error();
319 stats.error();
320
321 let minute2 = stats.snapshot() - minute1;
322
323 assert_eq!(
324 minute1.format_notime(),
325 "2 requests, avg latency: 0.15 ms, 1 retries"
326 );
327
328 assert_eq!(
329 minute2.format_notime(),
330 "1 requests, avg latency: 0.30 ms, 2 errors"
331 );
332 }
333}