1use crate::cli::output::{Output, RunRecord};
11use crate::{note, out};
12use scope::error::{Result, ScopeError};
13use scope::market::series::{self, Metric, Sample};
14use std::future::Future;
15use std::io::Write as _;
16use std::path::Path;
17use std::time::{Duration, Instant};
18use tokio::time::Instant as Clock;
20
21const CSV_TIME_FORMAT: &str = "%Y-%m-%dT%H:%M:%SZ";
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub struct Schedule {
27 pub every: Duration,
29 pub duration: Duration,
31}
32
33impl Schedule {
34 pub fn new(every: Duration, duration: Duration) -> Result<Self> {
36 if every.is_zero() {
37 return Err(ScopeError::InvalidInput(
38 "Interval must be positive".to_string(),
39 ));
40 }
41 Ok(Self { every, duration })
42 }
43
44 pub fn total(&self) -> u64 {
48 let every = self.every.as_millis().max(1);
49 let total = self.duration.as_millis().div_ceil(every).saturating_add(1);
50 u64::try_from(total).unwrap_or(u64::MAX)
51 }
52}
53
54#[derive(Debug, Clone)]
56pub struct RunOutput<T> {
57 pub sample: Sample,
59 pub snapshot: String,
61 pub value: T,
63}
64
65pub async fn run<T, F, Fut>(
73 schedule: &Schedule,
74 metrics: &[Metric],
75 csv: Option<&Path>,
76 out: &Output,
77 mut step: F,
78) -> Result<Option<T>>
79where
80 F: FnMut(u64) -> Fut,
81 Fut: Future<Output = Result<RunOutput<T>>>,
82{
83 let columns = series::csv_columns(metrics).map_err(|e| ScopeError::Other(e.to_string()))?;
84 if let Some(path) = csv {
85 std::fs::write(path, series::csv_header(&columns))
86 .map_err(|e| ScopeError::io_context(e, format!("cannot write {}", path.display())))?;
87 }
88 out.update_progress(|p| {
89 p.run = 0;
90 p.total = Some(schedule.total());
91 p.next_at = None;
92 });
93
94 let tui = out.series_hook().is_some();
95 let start = Clock::now();
96 let end = start + schedule.duration;
97 let mut slot = start;
98 let mut run: u64 = 1;
99 loop {
100 out.update_progress(|p| {
101 p.run = run;
102 p.next_at = None;
103 });
104 if !tui {
105 let ts = chrono::Utc::now().format("%H:%M:%S");
106 note!(out, "--- Run #{} at {} ---\n", run, ts);
107 }
108 let output = step(run).await?;
110 let at = chrono::Utc::now();
111 if let Some(path) = csv {
112 let time = at.format(CSV_TIME_FORMAT).to_string();
113 let row = series::csv_row(&columns, metrics, &time, run, &output.sample);
114 append(path, &row)?;
115 }
116 let record = RunRecord {
117 run,
118 at,
119 sample: output.sample,
120 snapshot: output.snapshot.clone(),
121 };
122 if !out.record_run(metrics, record) {
123 out!(out, "{}", output.snapshot)?;
124 }
125 let last = output.value;
126
127 if start.elapsed() >= schedule.duration {
128 note!(
129 out,
130 "\nCompleted {} run(s) over {}s.",
131 run,
132 schedule.duration.as_secs()
133 );
134 return Ok(Some(last));
135 }
136 run += 1;
137 slot = (slot + schedule.every).max(Clock::now()).min(end);
141 let pause = slot.saturating_duration_since(Clock::now());
142 out.update_progress(|p| p.next_at = Some(Instant::now() + pause));
143 tokio::time::sleep_until(slot).await;
144 }
145}
146
147fn append(path: &Path, row: &str) -> Result<()> {
148 let context = || format!("cannot append to {}", path.display());
149 let mut file = std::fs::OpenOptions::new()
150 .append(true)
151 .open(path)
152 .map_err(|e| ScopeError::io_context(e, context()))?;
153 file.write_all(row.as_bytes())
154 .map_err(|e| ScopeError::io_context(e, context()))
155}
156
157#[cfg(test)]
158mod tests {
159 use super::*;
160 use scope::market::HealthThresholds;
161 use scope::market::series::summary_metrics;
162
163 fn metrics() -> Vec<Metric> {
164 summary_metrics(&HealthThresholds::default())
165 }
166
167 fn output(run: u64, metrics: &[Metric]) -> RunOutput<u64> {
168 RunOutput {
169 sample: Sample {
170 values: vec![Some(run as f64); metrics.len()],
171 healthy: true,
172 failed_checks: vec![],
173 },
174 snapshot: format!("report {run}\n"),
175 value: run,
176 }
177 }
178
179 #[test]
180 fn test_total_counts_the_final_run_at_the_boundary() {
181 let s = |every, duration| {
184 Schedule::new(Duration::from_secs(every), Duration::from_secs(duration)).unwrap()
185 };
186 assert_eq!(s(2, 5).total(), 4);
187 assert_eq!(s(1, 2).total(), 3);
188 assert_eq!(s(2, 1).total(), 2);
189 }
190
191 #[test]
192 fn test_a_zero_interval_is_the_users_error() {
193 let err = Schedule::new(Duration::ZERO, Duration::from_secs(1)).unwrap_err();
194 assert!(matches!(err, ScopeError::InvalidInput(_)), "{err:?}");
195 }
196
197 #[tokio::test(start_paused = true)]
198 async fn test_cli_prints_each_report_and_returns_the_last_value() {
199 let m = metrics();
200 let (out, cap) = Output::capture();
201 let schedule = Schedule::new(Duration::from_secs(1), Duration::from_secs(2)).unwrap();
202 let last = run(&schedule, &m, None, &out, |n| {
203 let o = output(n, &m);
204 async move { Ok(o) }
205 })
206 .await
207 .unwrap();
208 assert_eq!(last, Some(3));
209 assert_eq!(cap.out(), "report 1\nreport 2\nreport 3\n");
210 assert!(cap.err().contains("--- Run #3 at"), "{}", cap.err());
211 assert!(
212 cap.err().contains("Completed 3 run(s) over 2s."),
213 "{}",
214 cap.err()
215 );
216 }
217
218 #[tokio::test(start_paused = true)]
219 async fn test_tui_gets_records_and_no_repeated_reports() {
220 let m = metrics();
221 let (out, cap) = Output::capture_merged();
222 let schedule = Schedule::new(Duration::from_secs(1), Duration::from_secs(1)).unwrap();
223 run(&schedule, &m, None, &out, |n| {
224 let o = output(n, &m);
225 async move { Ok(o) }
226 })
227 .await
228 .unwrap();
229 assert!(!cap.out().contains("report 1"), "{}", cap.out());
230 assert!(!cap.out().contains("--- Run #"), "{}", cap.out());
231 let series = out.series_hook().unwrap();
232 let series = series.lock().unwrap();
233 assert_eq!(series.metrics, m);
234 let runs: Vec<u64> = series.runs.iter().map(|r| r.run).collect();
235 assert_eq!(runs, vec![1, 2]);
236 assert_eq!(series.runs[1].snapshot, "report 2\n");
237 }
238
239 #[tokio::test(start_paused = true)]
240 async fn test_an_error_ends_the_repeat() {
241 let m = metrics();
242 let (out, _cap) = Output::capture_merged();
243 let schedule = Schedule::new(Duration::from_secs(1), Duration::from_secs(30)).unwrap();
244 let mut calls = 0;
245 let result: Result<Option<u64>> = run(&schedule, &m, None, &out, |n| {
246 calls += 1;
247 let o = output(n, &m);
248 async move {
249 if n == 2 {
250 Err(ScopeError::Network("venue down".into()))
251 } else {
252 Ok(o)
253 }
254 }
255 })
256 .await;
257 assert!(result.is_err());
258 assert_eq!(calls, 2, "no run after the error");
259 let progress = out.progress_hook().unwrap().lock().unwrap().clone();
260 assert_eq!(progress.run, 2);
261 assert_eq!(out.series_hook().unwrap().lock().unwrap().runs.len(), 1);
263 }
264
265 #[tokio::test(start_paused = true)]
266 async fn test_csv_gets_the_header_then_one_row_per_run() {
267 let m = metrics();
268 let dir = tempfile::tempdir().unwrap();
269 let path = dir.path().join("s.csv");
270 let (out, _cap) = Output::capture();
271 let schedule = Schedule::new(Duration::from_secs(1), Duration::from_secs(1)).unwrap();
272 run(&schedule, &m, Some(&path), &out, |n| {
273 let o = output(n, &m);
274 async move { Ok(o) }
275 })
276 .await
277 .unwrap();
278 let body = std::fs::read_to_string(&path).unwrap();
279 let lines: Vec<&str> = body.lines().collect();
280 let columns = series::csv_columns(&m).unwrap();
281 assert_eq!(format!("{}\n", lines[0]), series::csv_header(&columns));
282 assert_eq!(lines.len(), 3, "{body}");
283 assert_eq!(lines[2].split(',').nth(1), Some("2"));
284 }
285
286 async fn start_times(schedule: Schedule, cost: impl Fn(u64) -> Duration) -> Vec<u128> {
289 let m = metrics();
290 let (out, _cap) = Output::capture_merged();
291 let start = Clock::now();
292 let starts = std::sync::Mutex::new(Vec::new());
293 run(&schedule, &m, None, &out, |n| {
294 starts.lock().unwrap().push(start.elapsed().as_millis());
295 let o = output(n, &m);
296 let cost = cost(n);
297 async move {
298 tokio::time::sleep(cost).await;
299 Ok(o)
300 }
301 })
302 .await
303 .unwrap();
304 starts.into_inner().unwrap()
305 }
306
307 #[tokio::test(start_paused = true)]
308 async fn test_run_time_does_not_push_later_runs_back() {
309 let schedule = Schedule::new(Duration::from_secs(30), Duration::from_secs(90)).unwrap();
312 let starts = start_times(schedule, |_| Duration::from_millis(750)).await;
313 assert_eq!(starts, vec![0, 30_000, 60_000, 90_000]);
314 assert_eq!(starts.len() as u64, schedule.total());
315 }
316
317 #[tokio::test(start_paused = true)]
318 async fn test_an_overrun_moves_the_grid_without_a_burst() {
319 let schedule = Schedule::new(Duration::from_secs(1), Duration::from_secs(6)).unwrap();
323 let cost = |n| Duration::from_millis(if n == 1 { 3_500 } else { 100 });
324 let starts = start_times(schedule, cost).await;
325 assert_eq!(starts, vec![0, 3_500, 4_500, 5_500, 6_000]);
326 }
327
328 #[test]
329 fn test_the_series_ring_keeps_the_newest_runs() {
330 let m = metrics();
331 let (out, _cap) = Output::capture_merged();
332 let extra = 3;
333 for n in 1..=(crate::cli::output::MAX_SERIES_RUNS as u64 + extra) {
334 let o = output(n, &m);
335 out.record_run(
336 &m,
337 RunRecord {
338 run: n,
339 at: chrono::Utc::now(),
340 sample: o.sample,
341 snapshot: o.snapshot,
342 },
343 );
344 }
345 let series = out.series_hook().unwrap();
346 let series = series.lock().unwrap();
347 assert_eq!(series.runs.len(), crate::cli::output::MAX_SERIES_RUNS);
348 assert_eq!(series.dropped, extra);
349 assert_eq!(series.runs.front().unwrap().run, extra + 1);
350 }
351}