Skip to main content

seyk_core/
progress.rs

1//! Resource-level progress for one indexing pass. No persistent cache state.
2use 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/// Counts of discovered resources in the current (or last completed) pass.
10/// A mixed PDF is one `pdf` resource, regardless of its page locator kinds.
11#[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/// Snapshot of a pass, not a durable inventory. Counts reset on rediscovery.
19#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
20pub struct IndexProgress {
21    pub by_kind: BTreeMap<String, KindProgress>,
22    /// Remaining resource count / observed newly indexed resources per second.
23    /// Null until two new resources complete and at least 1 ms has elapsed;
24    /// cached reuse and skips never train the estimate. Null when finished.
25    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        // A provider repeating an ID must not double-count or replace its kind.
149        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); // reuse is not a sample
157        assert_eq!(tracker.snapshot_after(elapsed).eta_seconds, None);
158        tracker.complete(&resources[1].id, Completion::Skipped); // nor is a skip
159        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}