1use crate::storage::{KeyValueStore, KeyValueStoreExt, StorageResult};
4use serde::{Deserialize, Serialize};
5use std::{
6 collections::{BTreeMap, VecDeque},
7 sync::{Arc, Mutex, MutexGuard},
8 time::{Duration, Instant},
9};
10use tokio::sync::mpsc;
11
12pub const STATISTICS_PERSIST_KEY: &str = "SDK_CRAWLER_STATISTICS_0";
14
15const DEFAULT_WINDOW: Duration = Duration::from_secs(60);
16const DEFAULT_UPDATE_INTERVAL: Duration = Duration::from_secs(1);
17const ERROR_KEY_MAX_CHARS: usize = 200;
18
19#[derive(Clone)]
21pub struct StatisticsHandle {
22 inner: Arc<StatisticsInner>,
23}
24
25struct StatisticsInner {
26 state: Mutex<State>,
27 subscribers: Mutex<Vec<mpsc::Sender<StatisticsSnapshot>>>,
28 window: Duration,
29}
30
31struct State {
32 requests_finished: u64,
33 requests_failed: u64,
34 requests_retries: u64,
35 request_duration_sum: Duration,
36 request_min_duration: Duration,
37 request_max_duration: Duration,
38 status_codes: BTreeMap<u16, u64>,
39 retry_histogram: Vec<u64>,
40 errors: BTreeMap<String, u64>,
41 retry_errors: BTreeMap<String, u64>,
42 accumulated_runtime: Duration,
43 started_at: Option<Instant>,
44 recent_events: VecDeque<(Instant, bool)>,
45}
46
47impl Default for State {
48 fn default() -> Self {
49 Self {
50 requests_finished: 0,
51 requests_failed: 0,
52 requests_retries: 0,
53 request_duration_sum: Duration::ZERO,
54 request_min_duration: Duration::ZERO,
55 request_max_duration: Duration::ZERO,
56 status_codes: BTreeMap::new(),
57 retry_histogram: Vec::new(),
58 errors: BTreeMap::new(),
59 retry_errors: BTreeMap::new(),
60 accumulated_runtime: Duration::ZERO,
61 started_at: None,
62 recent_events: VecDeque::new(),
63 }
64 }
65}
66
67#[derive(Debug, Clone, PartialEq, Serialize)]
69pub struct StatisticsSnapshot {
70 pub requests_finished: u64,
72 pub requests_failed: u64,
74 pub requests_retries: u64,
76 pub requests_finished_per_minute: f64,
78 pub requests_failed_per_minute: f64,
80 pub request_avg_duration: Duration,
82 pub request_min_duration: Duration,
84 pub request_max_duration: Duration,
86 pub status_codes: BTreeMap<u16, u64>,
88 pub crawler_runtime: Duration,
90 pub retry_histogram: Vec<u64>,
92 pub errors: BTreeMap<String, u64>,
94 pub retry_errors: BTreeMap<String, u64>,
96}
97
98#[derive(Debug, Clone, PartialEq, Serialize)]
100#[must_use = "final statistics summarize the completed crawler run"]
101pub struct FinalStatistics {
102 pub requests_finished: u64,
104 pub requests_failed: u64,
106 pub requests_retries: u64,
108 pub requests_finished_per_minute: f64,
110 pub requests_failed_per_minute: f64,
112 pub request_avg_duration: Duration,
114 pub request_min_duration: Duration,
116 pub request_max_duration: Duration,
118 pub status_codes: BTreeMap<u16, u64>,
120 pub crawler_runtime: Duration,
122 pub retry_histogram: Vec<u64>,
124 pub errors: BTreeMap<String, u64>,
126 pub retry_errors: BTreeMap<String, u64>,
128}
129
130impl From<StatisticsSnapshot> for FinalStatistics {
131 fn from(snapshot: StatisticsSnapshot) -> Self {
132 Self {
133 requests_finished: snapshot.requests_finished,
134 requests_failed: snapshot.requests_failed,
135 requests_retries: snapshot.requests_retries,
136 requests_finished_per_minute: snapshot.requests_finished_per_minute,
137 requests_failed_per_minute: snapshot.requests_failed_per_minute,
138 request_avg_duration: snapshot.request_avg_duration,
139 request_min_duration: snapshot.request_min_duration,
140 request_max_duration: snapshot.request_max_duration,
141 status_codes: snapshot.status_codes,
142 crawler_runtime: snapshot.crawler_runtime,
143 retry_histogram: snapshot.retry_histogram,
144 errors: snapshot.errors,
145 retry_errors: snapshot.retry_errors,
146 }
147 }
148}
149
150#[derive(Serialize, Deserialize)]
151struct PersistedStatistics {
152 requests_finished: u64,
153 requests_failed: u64,
154 requests_retries: u64,
155 request_duration_sum: Duration,
156 request_min_duration: Duration,
157 request_max_duration: Duration,
158 status_codes: BTreeMap<u16, u64>,
159 retry_histogram: Vec<u64>,
160 errors: BTreeMap<String, u64>,
161 retry_errors: BTreeMap<String, u64>,
162 accumulated_runtime: Duration,
163}
164
165impl StatisticsHandle {
166 #[must_use]
168 pub fn new() -> Self {
169 Self::with_window(DEFAULT_WINDOW)
170 }
171
172 pub(crate) fn with_window(window: Duration) -> Self {
174 Self {
175 inner: Arc::new(StatisticsInner {
176 state: Mutex::new(State::default()),
177 subscribers: Mutex::new(Vec::new()),
178 window,
179 }),
180 }
181 }
182
183 pub fn record_finished(&self, duration: Duration, status_code: Option<u16>, retry_count: u32) {
187 let now = Instant::now();
188 let mut state = self.lock();
189 state.requests_finished += 1;
190 state.record_duration(duration);
191 state.record_terminal_retry_count(retry_count);
192 if let Some(status_code) = status_code {
193 *state.status_codes.entry(status_code).or_default() += 1;
194 }
195 state.recent_events.push_back((now, true));
196 prune_window(&mut state, now, self.inner.window);
197 drop(state);
198 self.emit_snapshot();
199 }
200
201 pub fn record_failed(&self, duration: Duration, error_key: &str, retry_count: u32) {
205 let now = Instant::now();
206 let mut state = self.lock();
207 state.requests_failed += 1;
208 state.record_duration(duration);
209 state.record_terminal_retry_count(retry_count);
210 *state
211 .errors
212 .entry(normalize_error_key(error_key))
213 .or_default() += 1;
214 state.recent_events.push_back((now, false));
215 prune_window(&mut state, now, self.inner.window);
216 drop(state);
217 self.emit_snapshot();
218 }
219
220 pub fn record_retry(&self, error_key: &str) {
222 let now = Instant::now();
223 let mut state = self.lock();
224 state.requests_retries += 1;
225 *state
226 .retry_errors
227 .entry(normalize_error_key(error_key))
228 .or_default() += 1;
229 prune_window(&mut state, now, self.inner.window);
230 drop(state);
231 self.emit_snapshot();
232 }
233
234 pub fn mark_run_started(&self) {
236 let mut state = self.lock();
237 if state.started_at.is_none() {
238 state.started_at = Some(Instant::now());
239 }
240 }
241
242 pub fn mark_run_stopped(&self) {
244 let mut state = self.lock();
245 if let Some(started_at) = state.started_at.take() {
246 state.accumulated_runtime += started_at.elapsed();
247 }
248 }
249
250 #[must_use]
252 pub fn snapshot(&self) -> StatisticsSnapshot {
253 let now = Instant::now();
254 let mut state = self.lock();
255 prune_window(&mut state, now, self.inner.window);
256 let (finished_in_window, failed_in_window) = state.recent_events.iter().fold(
257 (0_u64, 0_u64),
258 |(finished, failed), (_, succeeded)| {
259 if *succeeded {
260 (finished + 1, failed)
261 } else {
262 (finished, failed + 1)
263 }
264 },
265 );
266 let rate_factor = if self.inner.window.is_zero() {
267 0.0
268 } else {
269 60.0 / self.inner.window.as_secs_f64()
270 };
271 let completed = state.requests_finished + state.requests_failed;
272 StatisticsSnapshot {
273 requests_finished: state.requests_finished,
274 requests_failed: state.requests_failed,
275 requests_retries: state.requests_retries,
276 requests_finished_per_minute: finished_in_window as f64 * rate_factor,
277 requests_failed_per_minute: failed_in_window as f64 * rate_factor,
278 request_avg_duration: if completed == 0 {
279 Duration::ZERO
280 } else {
281 duration_average(state.request_duration_sum, completed)
282 },
283 request_min_duration: state.request_min_duration,
284 request_max_duration: state.request_max_duration,
285 status_codes: state.status_codes.clone(),
286 crawler_runtime: state.accumulated_runtime
287 + state.started_at.map_or(Duration::ZERO, |start| now - start),
288 retry_histogram: state.retry_histogram.clone(),
289 errors: state.errors.clone(),
290 retry_errors: state.retry_errors.clone(),
291 }
292 }
293
294 #[must_use]
301 pub fn subscribe(&self) -> mpsc::Receiver<StatisticsSnapshot> {
302 let (sender, receiver) = mpsc::channel(16);
303 self.inner
304 .subscribers
305 .lock()
306 .expect("statistics subscribers mutex poisoned")
307 .push(sender.clone());
308
309 let statistics = self.clone();
310 let update_interval = if self.inner.window.is_zero() {
311 DEFAULT_UPDATE_INTERVAL
312 } else {
313 self.inner.window.min(DEFAULT_UPDATE_INTERVAL)
314 };
315 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
316 runtime.spawn(async move {
317 let mut interval = tokio::time::interval(update_interval);
318 interval.tick().await;
319 loop {
320 interval.tick().await;
321 if sender.send(statistics.snapshot()).await.is_err() {
322 statistics
323 .inner
324 .subscribers
325 .lock()
326 .expect("statistics subscribers mutex poisoned")
327 .retain(|subscriber| !subscriber.is_closed());
328 break;
329 }
330 }
331 });
332 }
333 receiver
334 }
335
336 pub fn finalize(&self) -> FinalStatistics {
338 self.snapshot().into()
339 }
340
341 pub async fn persist(&self, kvs: &dyn KeyValueStore) -> StorageResult<()> {
343 let persisted = {
344 let now = Instant::now();
345 let state = self.lock();
346 PersistedStatistics {
347 requests_finished: state.requests_finished,
348 requests_failed: state.requests_failed,
349 requests_retries: state.requests_retries,
350 request_duration_sum: state.request_duration_sum,
351 request_min_duration: state.request_min_duration,
352 request_max_duration: state.request_max_duration,
353 status_codes: state.status_codes.clone(),
354 retry_histogram: state.retry_histogram.clone(),
355 errors: state.errors.clone(),
356 retry_errors: state.retry_errors.clone(),
357 accumulated_runtime: state.accumulated_runtime
358 + state.started_at.map_or(Duration::ZERO, |start| now - start),
359 }
360 };
361 kvs.set(STATISTICS_PERSIST_KEY, &persisted).await
362 }
363
364 pub async fn restore(&self, kvs: &dyn KeyValueStore) -> StorageResult<bool> {
368 let Some(persisted) = kvs
369 .get::<PersistedStatistics>(STATISTICS_PERSIST_KEY)
370 .await?
371 else {
372 return Ok(false);
373 };
374 let mut state = self.lock();
375 *state = State {
376 requests_finished: persisted.requests_finished,
377 requests_failed: persisted.requests_failed,
378 requests_retries: persisted.requests_retries,
379 request_duration_sum: persisted.request_duration_sum,
380 request_min_duration: persisted.request_min_duration,
381 request_max_duration: persisted.request_max_duration,
382 status_codes: persisted.status_codes,
383 retry_histogram: persisted.retry_histogram,
384 errors: persisted.errors,
385 retry_errors: persisted.retry_errors,
386 accumulated_runtime: persisted.accumulated_runtime,
387 started_at: None,
388 recent_events: VecDeque::new(),
389 };
390 Ok(true)
391 }
392
393 fn lock(&self) -> MutexGuard<'_, State> {
394 self.inner.state.lock().expect("statistics mutex poisoned")
395 }
396
397 fn emit_snapshot(&self) {
398 let snapshot = self.snapshot();
399 self.inner
400 .subscribers
401 .lock()
402 .expect("statistics subscribers mutex poisoned")
403 .retain(|subscriber| {
404 !matches!(
405 subscriber.try_send(snapshot.clone()),
406 Err(mpsc::error::TrySendError::Closed(_))
407 )
408 });
409 }
410}
411
412impl Default for StatisticsHandle {
413 fn default() -> Self {
414 Self::new()
415 }
416}
417
418impl State {
419 fn record_duration(&mut self, duration: Duration) {
420 let first = self.requests_finished + self.requests_failed == 1;
421 self.request_duration_sum += duration;
422 if first || duration < self.request_min_duration {
423 self.request_min_duration = duration;
424 }
425 if first || duration > self.request_max_duration {
426 self.request_max_duration = duration;
427 }
428 }
429
430 fn record_terminal_retry_count(&mut self, retry_count: u32) {
431 let index = retry_count as usize;
432 if self.retry_histogram.len() <= index {
433 self.retry_histogram.resize(index + 1, 0);
434 }
435 self.retry_histogram[index] += 1;
436 }
437}
438
439fn normalize_error_key(error_key: &str) -> String {
440 let first_line = error_key.lines().next().unwrap_or(error_key);
441 let mut normalized = String::new();
442
443 for token in first_line.split_whitespace() {
444 if !normalized.is_empty() {
445 normalized.push(' ');
446 }
447
448 if token.starts_with("http://") || token.starts_with("https://") {
449 normalized.push_str("<url>");
450 } else if is_uuid_token(token) {
451 normalized.push_str("<uuid>");
452 } else {
453 let mut previous_was_digit = false;
454 for character in token.chars() {
455 if character.is_ascii_digit() {
456 if !previous_was_digit {
457 normalized.push('#');
458 }
459 previous_was_digit = true;
460 } else {
461 normalized.push(character);
462 previous_was_digit = false;
463 }
464 }
465 }
466 }
467
468 normalized.chars().take(ERROR_KEY_MAX_CHARS).collect()
469}
470
471fn is_uuid_token(token: &str) -> bool {
472 let expected_lengths = [8, 4, 4, 4, 12];
473 let mut parts = token.split('-');
474
475 for expected_length in expected_lengths {
476 let Some(part) = parts.next() else {
477 return false;
478 };
479 if part.len() != expected_length || !part.bytes().all(|byte| byte.is_ascii_hexdigit()) {
480 return false;
481 }
482 }
483
484 parts.next().is_none()
485}
486
487fn duration_average(total: Duration, count: u64) -> Duration {
488 let seconds = total.as_secs() / count;
489 let remaining_seconds = total.as_secs() % count;
490 let remaining_nanos =
491 u128::from(remaining_seconds) * 1_000_000_000 + u128::from(total.subsec_nanos());
492 Duration::new(seconds, (remaining_nanos / u128::from(count)) as u32)
493}
494
495fn prune_window(state: &mut State, now: Instant, window: Duration) {
496 while state
497 .recent_events
498 .front()
499 .is_some_and(|(recorded_at, _)| now.duration_since(*recorded_at) > window)
500 {
501 state.recent_events.pop_front();
502 }
503}
504
505#[cfg(test)]
506mod tests {
507 use super::*;
508
509 #[test]
510 fn counters() {
511 let statistics = StatisticsHandle::new();
512 statistics.record_finished(Duration::from_millis(10), None, 0);
513 statistics.record_finished(Duration::from_millis(20), None, 1);
514 statistics.record_finished(Duration::from_millis(30), None, 1);
515 statistics.record_failed(Duration::from_millis(40), "non-retryable: boom", 2);
516 for _ in 0..3 {
517 statistics.record_retry("retryable: flaky");
518 }
519
520 let snapshot = statistics.snapshot();
521 assert_eq!(snapshot.requests_finished, 3);
522 assert_eq!(snapshot.requests_failed, 1);
523 assert_eq!(snapshot.requests_retries, 3);
524 assert_eq!(snapshot.retry_histogram, vec![1, 2, 1]);
525 assert_eq!(
526 snapshot.errors,
527 BTreeMap::from([("non-retryable: boom".to_owned(), 1)])
528 );
529 assert_eq!(
530 snapshot.retry_errors,
531 BTreeMap::from([("retryable: flaky".to_owned(), 3)])
532 );
533 assert_eq!(snapshot.request_min_duration, Duration::from_millis(10));
534 assert_eq!(snapshot.request_max_duration, Duration::from_millis(40));
535 assert_eq!(snapshot.request_avg_duration, Duration::from_millis(25));
536 }
537
538 #[test]
539 fn normalize_error_key_collapses_digits() {
540 assert_eq!(
541 normalize_error_key("non-retryable: boom 42"),
542 "non-retryable: boom #"
543 );
544 assert_eq!(
545 normalize_error_key("non-retryable: boom 999"),
546 "non-retryable: boom #"
547 );
548 assert_eq!(
549 normalize_error_key("timeout after 5031ms"),
550 "timeout after #ms"
551 );
552 }
553
554 #[test]
555 fn normalize_error_key_masks_urls() {
556 assert_eq!(
557 normalize_error_key("non-retryable: fetch https://a.com/1 failed"),
558 "non-retryable: fetch <url> failed"
559 );
560 assert_eq!(
561 normalize_error_key("non-retryable: fetch https://b.com/2 failed"),
562 "non-retryable: fetch <url> failed"
563 );
564 }
565
566 #[test]
567 fn normalize_error_key_masks_uuid() {
568 assert_eq!(
569 normalize_error_key("session: id 550e8400-e29b-41d4-a716-446655440000 failed"),
570 "session: id <uuid> failed"
571 );
572 }
573
574 #[test]
575 fn normalize_error_key_only_first_line() {
576 assert_eq!(
577 normalize_error_key("retryable: boom\ncaused by: id 42"),
578 "retryable: boom"
579 );
580 }
581
582 #[tokio::test]
583 async fn similar_failures_group_into_one_bucket() {
584 let statistics = StatisticsHandle::new();
585 statistics.record_failed(
586 Duration::from_millis(10),
587 "non-retryable: fetch https://x/1 failed",
588 0,
589 );
590 statistics.record_failed(
591 Duration::from_millis(10),
592 "non-retryable: fetch https://y/2 failed",
593 0,
594 );
595
596 let snapshot = statistics.snapshot();
597 assert_eq!(snapshot.errors.len(), 1);
598 assert_eq!(snapshot.errors.values().next(), Some(&2));
599
600 statistics.record_failed(
601 Duration::from_millis(10),
602 "session: fetch https://z/3 failed",
603 0,
604 );
605 let snapshot = statistics.snapshot();
606 assert_eq!(snapshot.errors.len(), 2);
607 assert!(
608 snapshot
609 .errors
610 .contains_key("non-retryable: fetch <url> failed")
611 );
612 assert!(snapshot.errors.contains_key("session: fetch <url> failed"));
613 }
614
615 #[test]
616 fn subscribe_outside_tokio_runtime_does_not_panic() {
617 let statistics = StatisticsHandle::new();
618 let mut snapshots = statistics.subscribe();
619
620 statistics.record_finished(Duration::ZERO, None, 0);
621
622 assert_eq!(
623 snapshots
624 .try_recv()
625 .expect("record operation should emit a snapshot")
626 .requests_finished,
627 1
628 );
629 }
630
631 #[tokio::test]
632 async fn sliding_window() {
633 let statistics = StatisticsHandle::with_window(Duration::from_millis(50));
634 for _ in 0..3 {
635 statistics.record_finished(Duration::ZERO, None, 0);
636 }
637 tokio::time::sleep(Duration::from_millis(80)).await;
638 statistics.record_finished(Duration::ZERO, None, 0);
639
640 let snapshot = statistics.snapshot();
641 assert_eq!(snapshot.requests_finished_per_minute, 1_200.0);
642 assert_eq!(snapshot.requests_finished, 4);
643 }
644
645 #[tokio::test]
646 async fn subscriber_receives_recorded_state() {
647 let statistics = StatisticsHandle::new();
648 let mut snapshots = statistics.subscribe();
649
650 statistics.record_finished(Duration::from_millis(10), Some(200), 0);
651
652 let snapshot = snapshots.recv().await.expect("snapshot channel closed");
653 assert_eq!(snapshot.requests_finished, 1);
654 assert_eq!(snapshot.status_codes, BTreeMap::from([(200, 1)]));
655 }
656
657 #[tokio::test]
658 async fn subscriber_receives_periodic_updates_without_new_events() {
659 let statistics = StatisticsHandle::with_window(Duration::from_millis(20));
660 statistics.mark_run_started();
661 let mut snapshots = statistics.subscribe();
662 statistics.record_finished(Duration::ZERO, None, 0);
663
664 let recorded = snapshots.recv().await.expect("snapshot channel closed");
665 assert!(recorded.requests_finished_per_minute > 0.0);
666 let recorded_runtime = recorded.crawler_runtime;
667
668 tokio::time::sleep(Duration::from_millis(50)).await;
669 let decayed = loop {
670 let snapshot = snapshots.recv().await.expect("snapshot channel closed");
671 if snapshot.requests_finished_per_minute == 0.0 {
672 break snapshot;
673 }
674 };
675 assert!(decayed.crawler_runtime > recorded_runtime);
676 }
677
678 #[tokio::test]
679 async fn runtime_accumulates_across_start_stop_and_is_monotonic_while_running() {
680 let statistics = StatisticsHandle::new();
681 statistics.mark_run_started();
682 tokio::time::sleep(Duration::from_millis(5)).await;
683 let first = statistics.snapshot().crawler_runtime;
684 tokio::time::sleep(Duration::from_millis(5)).await;
685 let second = statistics.snapshot().crawler_runtime;
686 assert!(second >= first);
687 statistics.mark_run_stopped();
688 let stopped = statistics.snapshot().crawler_runtime;
689 statistics.mark_run_stopped();
690 statistics.mark_run_started();
691 tokio::time::sleep(Duration::from_millis(5)).await;
692 statistics.mark_run_stopped();
693 assert!(statistics.snapshot().crawler_runtime > stopped);
694 }
695
696 #[test]
697 fn status_codes() {
698 let statistics = StatisticsHandle::new();
699 statistics.record_finished(Duration::ZERO, Some(200), 0);
700 statistics.record_finished(Duration::ZERO, Some(200), 0);
701 statistics.record_finished(Duration::ZERO, Some(404), 0);
702 assert_eq!(
703 statistics.snapshot().status_codes,
704 BTreeMap::from([(200, 2), (404, 1)])
705 );
706 }
707}