1use crate::discovery::ResourceKind;
3use crate::{ProgressEvent, ProgressFn, ProvidedResource, ResourceId};
4use serde::{Deserialize, Serialize};
5use std::collections::{BTreeMap, HashMap};
6use std::sync::Mutex;
7use std::time::{Duration, Instant};
8
9#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
12pub struct KindProgress {
13 pub indexed: usize,
14 pub waiting: usize,
15 pub skipped: usize,
16}
17
18#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
20pub struct IndexProgress {
21 pub by_kind: BTreeMap<String, KindProgress>,
22 pub eta_seconds: Option<u64>,
26}
27
28pub(crate) enum Completion {
29 Skipped,
30 Reused,
31 NewlyIndexed,
32}
33
34pub(crate) struct PassProgress<'a> {
35 state: Mutex<PassState>,
36 callback: Option<ProgressFn<'a>>,
37 started: Instant,
38}
39struct PassState {
40 pending: HashMap<ResourceId, &'static str>,
41 snapshot: IndexProgress,
42 updated: usize,
43}
44impl<'a> PassProgress<'a> {
45 pub(crate) fn new(resources: &[ProvidedResource], callback: Option<ProgressFn<'a>>) -> Self {
46 let mut state = PassState {
47 pending: HashMap::new(),
48 snapshot: IndexProgress::default(),
49 updated: 0,
50 };
51 for r in resources {
52 let kind = match r.kind {
53 ResourceKind::Text(_) => "text",
54 ResourceKind::Pdf => "pdf",
55 ResourceKind::Image => "image",
56 ResourceKind::Audio => "audio",
57 };
58 if let std::collections::hash_map::Entry::Vacant(entry) =
59 state.pending.entry(r.id.clone())
60 {
61 entry.insert(kind);
62 state
63 .snapshot
64 .by_kind
65 .entry(kind.into())
66 .or_default()
67 .waiting += 1;
68 }
69 }
70 let tracker = Self {
71 state: Mutex::new(state),
72 callback,
73 started: Instant::now(),
74 };
75 tracker.emit();
76 tracker
77 }
78 fn emit(&self) {
79 if let Some(callback) = self.callback {
80 let progress = self.snapshot();
81 callback(ProgressEvent::IndexProgress { progress });
82 }
83 }
84 pub(crate) fn complete(&self, id: &ResourceId, completion: Completion) {
85 {
86 let mut state = self.state.lock().unwrap();
87 let Some(kind) = state.pending.remove(id) else {
88 return;
89 };
90 let counts = state.snapshot.by_kind.get_mut(kind).unwrap();
91 counts.waiting -= 1;
92 match completion {
93 Completion::Skipped => counts.skipped += 1,
94 Completion::Reused => counts.indexed += 1,
95 Completion::NewlyIndexed => {
96 counts.indexed += 1;
97 state.updated += 1;
98 }
99 }
100 }
101 self.emit();
102 }
103 pub(crate) fn snapshot(&self) -> IndexProgress {
104 self.snapshot_after(self.started.elapsed())
105 }
106 fn snapshot_after(&self, elapsed: Duration) -> IndexProgress {
107 let state = self.state.lock().unwrap();
108 let mut result = state.snapshot.clone();
109 result.eta_seconds = estimate(state.pending.len(), state.updated, elapsed);
110 result
111 }
112}
113fn estimate(waiting: usize, updated: usize, elapsed: Duration) -> Option<u64> {
114 if waiting == 0 || updated < 2 || elapsed < Duration::from_millis(1) {
115 return None;
116 }
117 Some((waiting as f64 * elapsed.as_secs_f64() / updated as f64).ceil() as u64)
118}
119#[cfg(test)]
120mod tests {
121 use super::*;
122 #[test]
123 fn eta_requires_observations_and_resets_at_completion() {
124 assert_eq!(estimate(10, 0, Duration::from_secs(10)), None);
125 assert_eq!(estimate(10, 1, Duration::from_secs(10)), None);
126 assert_eq!(estimate(10, 2, Duration::ZERO), None);
127 assert_eq!(estimate(10, 2, Duration::from_secs(4)), Some(20));
128 assert_eq!(estimate(6, 4, Duration::from_secs(6)), Some(9));
129 assert_eq!(estimate(0, 4, Duration::from_secs(6)), None);
130 }
131}
132
133#[cfg(test)]
134mod lifecycle_tests {
135 use super::*;
136 #[test]
137 fn eta_tracks_real_completion_transitions_with_a_controlled_clock() {
138 let mut resources: Vec<_> = (0..5)
139 .map(|i| ProvidedResource {
140 id: format!("mem://test/{i}").parse().unwrap(),
141 kind: ResourceKind::Text(crate::segment::TextFormat::Plain),
142 size_bytes: 1,
143 modified_ns: None,
144 version: Default::default(),
145 title: None,
146 })
147 .collect();
148 let mut duplicate = resources[0].clone();
150 duplicate.kind = ResourceKind::Image;
151 resources.push(duplicate);
152 let tracker = PassProgress::new(&resources, None);
153 assert_eq!(tracker.snapshot().by_kind.len(), 1);
154 let elapsed = Duration::from_secs(4);
155 assert_eq!(tracker.snapshot_after(elapsed).eta_seconds, None);
156 tracker.complete(&resources[0].id, Completion::Reused); assert_eq!(tracker.snapshot_after(elapsed).eta_seconds, None);
158 tracker.complete(&resources[1].id, Completion::Skipped); assert_eq!(tracker.snapshot_after(elapsed).eta_seconds, None);
160 let before = tracker.snapshot_after(elapsed);
161 tracker.complete(&resources[0].id, Completion::NewlyIndexed);
162 tracker.complete(&resources[1].id, Completion::NewlyIndexed);
163 assert_eq!(tracker.snapshot_after(elapsed), before);
164 tracker.complete(&resources[2].id, Completion::NewlyIndexed);
165 assert_eq!(tracker.snapshot_after(elapsed).eta_seconds, None);
166 tracker.complete(&resources[3].id, Completion::NewlyIndexed);
167 assert_eq!(tracker.snapshot_after(elapsed).eta_seconds, Some(2));
168 let before = tracker.snapshot_after(elapsed);
169 tracker.complete(&resources[3].id, Completion::Skipped);
170 tracker.complete(&resources[2].id, Completion::Reused);
171 assert_eq!(tracker.snapshot_after(elapsed), before);
172 tracker.complete(&resources[4].id, Completion::NewlyIndexed);
173 assert_eq!(tracker.snapshot_after(elapsed).eta_seconds, None);
174 let counts = &tracker.snapshot_after(elapsed).by_kind["text"];
175 assert_eq!((counts.indexed, counts.waiting, counts.skipped), (4, 0, 1));
176 }
177}