projectatlas_core/
index_work.rs

1//! Cooperative cancellation and failure contracts for bounded index work.
2
3use std::fmt;
4use std::sync::{
5    Arc,
6    atomic::{AtomicBool, AtomicU64, Ordering},
7};
8use std::time::{Duration, Instant};
9use thiserror::Error;
10
11/// Cloneable cancellation signal shared by one indexing operation.
12#[derive(Clone, Debug, Default)]
13pub struct IndexCancellation {
14    /// Atomic flag observed by every worker participating in the operation.
15    cancelled: Arc<AtomicBool>,
16}
17
18impl IndexCancellation {
19    /// Create an active cancellation signal.
20    #[must_use]
21    pub fn new() -> Self {
22        Self::default()
23    }
24
25    /// Request cooperative cancellation.
26    pub fn cancel(&self) {
27        self.cancelled.store(true, Ordering::Relaxed);
28    }
29
30    /// Return whether cancellation was requested.
31    #[must_use]
32    pub fn is_cancelled(&self) -> bool {
33        self.cancelled.load(Ordering::Relaxed)
34    }
35}
36
37/// Closed stages at which bounded indexing can stop.
38#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub enum IndexWorkStage {
40    /// Repository entries are being discovered and filtered.
41    RepositoryTraversal,
42    /// Source metadata is being read and classified.
43    SourceMetadata,
44    /// Exact source bytes are being hashed.
45    SourceHash,
46    /// Source text is being staged for lexical and structural indexing.
47    TextIndex,
48    /// Source symbols and relationships are being parsed.
49    SymbolParsing,
50    /// A completed repository scan is being finalized for its caller.
51    ScanFinalization,
52    /// Validated staged index data is being published.
53    Publication,
54}
55
56impl fmt::Display for IndexWorkStage {
57    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
58        formatter.write_str(match self {
59            Self::RepositoryTraversal => "repository_traversal",
60            Self::SourceMetadata => "source_metadata",
61            Self::SourceHash => "source_hash",
62            Self::TextIndex => "text_index",
63            Self::SymbolParsing => "symbol_parsing",
64            Self::ScanFinalization => "scan_finalization",
65            Self::Publication => "publication",
66        })
67    }
68}
69
70/// Closed resources governed by indexing limits.
71#[derive(Clone, Copy, Debug, Eq, PartialEq)]
72pub enum IndexWorkResource {
73    /// Repository entries considered by a scanner.
74    Entries,
75    /// Total source bytes admitted for exact content hashing.
76    SourceBytes,
77    /// Total UTF-8 source bytes retained for one text-index publication.
78    TextBytes,
79    /// Authored purpose and publication-input bytes inspected by one refresh.
80    PurposeBytes,
81    /// Normalized legacy purpose records admitted by one refresh.
82    PurposeRecords,
83    /// Symbol parse results retained before sequential persistence.
84    SymbolJobs,
85    /// Symbol rows admitted by one index publication.
86    SymbolRows,
87    /// Relation rows admitted by one index publication.
88    RelationRows,
89    /// Retained parser-output string bytes admitted by one index publication.
90    OutputBytes,
91    /// Parallel workers used by one indexing operation.
92    Workers,
93}
94
95impl fmt::Display for IndexWorkResource {
96    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
97        formatter.write_str(match self {
98            Self::Entries => "entries",
99            Self::SourceBytes => "source_bytes",
100            Self::TextBytes => "text_bytes",
101            Self::PurposeBytes => "purpose_bytes",
102            Self::PurposeRecords => "purpose_records",
103            Self::SymbolJobs => "symbol_jobs",
104            Self::SymbolRows => "symbol_rows",
105            Self::RelationRows => "relation_rows",
106            Self::OutputBytes => "output_bytes",
107            Self::Workers => "workers",
108        })
109    }
110}
111
112/// Typed cooperative-stop failures returned by bounded indexing work.
113#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
114pub enum IndexWorkFailure {
115    /// The caller requested cancellation.
116    #[error("index work was canceled during {stage}")]
117    Cancelled {
118        /// Stage observing cancellation.
119        stage: IndexWorkStage,
120    },
121    /// The operation reached its shared deadline.
122    #[error("index work deadline was reached during {stage}")]
123    DeadlineExceeded {
124        /// Stage observing the elapsed deadline.
125        stage: IndexWorkStage,
126    },
127    /// Work would exceed one declared resource limit.
128    #[error(
129        "index work exceeded the {resource} limit during {stage}: observed {observed}, limit {limit}"
130    )]
131    ResourceLimitExceeded {
132        /// Stage observing the exhausted resource.
133        stage: IndexWorkStage,
134        /// Resource whose limit was reached.
135        resource: IndexWorkResource,
136        /// Configured inclusive maximum.
137        limit: u64,
138        /// First observed value beyond the maximum.
139        observed: u64,
140    },
141}
142
143impl IndexWorkFailure {
144    /// Construct a typed resource-limit failure.
145    #[must_use]
146    pub const fn resource_limit(
147        stage: IndexWorkStage,
148        resource: IndexWorkResource,
149        limit: u64,
150        observed: u64,
151    ) -> Self {
152        Self::ResourceLimitExceeded {
153            stage,
154            resource,
155            limit,
156            observed,
157        }
158    }
159}
160
161/// Shared cancellation and deadline boundary for one indexing operation.
162#[derive(Clone, Debug)]
163pub struct IndexWorkControl {
164    /// Cooperative cancellation signal shared across operation workers.
165    cancellation: IndexCancellation,
166    /// Common operation start used to derive and report its deadline.
167    started_at: Instant,
168    /// Optional absolute deadline observed by every operation worker.
169    deadline: Option<Instant>,
170    /// Optional maximum worker count shared by every operation stage.
171    worker_ceiling: Option<usize>,
172    /// Authored-purpose bytes consumed by this operation and every clone.
173    purpose_bytes: Arc<AtomicU64>,
174}
175
176impl IndexWorkControl {
177    /// Create a work boundary with an optional timeout from one shared start.
178    #[must_use]
179    pub fn new(cancellation: IndexCancellation, timeout: Option<Duration>) -> Self {
180        let started_at = Instant::now();
181        // An unrepresentable deadline fails closed instead of silently removing the bound.
182        let deadline = timeout.map(|timeout| started_at.checked_add(timeout).unwrap_or(started_at));
183        Self {
184            cancellation,
185            started_at,
186            deadline,
187            worker_ceiling: None,
188            purpose_bytes: Arc::new(AtomicU64::new(0)),
189        }
190    }
191
192    /// Create a work boundary with a caller-selected absolute deadline.
193    #[must_use]
194    pub fn with_deadline(cancellation: IndexCancellation, deadline: Instant) -> Self {
195        Self {
196            cancellation,
197            started_at: Instant::now(),
198            deadline: Some(deadline),
199            worker_ceiling: None,
200            purpose_bytes: Arc::new(AtomicU64::new(0)),
201        }
202    }
203
204    /// Return the common start instant for this operation.
205    #[must_use]
206    pub fn started_at(&self) -> Instant {
207        self.started_at
208    }
209
210    /// Return the common absolute deadline, when configured.
211    #[must_use]
212    pub fn deadline(&self) -> Option<Instant> {
213        self.deadline
214    }
215
216    /// Return the maximum workers available to each operation stage, when bounded.
217    #[must_use]
218    pub fn worker_ceiling(&self) -> Option<usize> {
219        self.worker_ceiling
220    }
221
222    /// Clone this boundary while applying a worker ceiling to every operation stage.
223    #[must_use]
224    pub fn with_worker_ceiling(&self, max_workers: usize) -> Self {
225        let ceiling = max_workers.max(1);
226        Self {
227            cancellation: self.cancellation.clone(),
228            started_at: self.started_at,
229            deadline: self.deadline,
230            worker_ceiling: Some(
231                self.worker_ceiling
232                    .map_or(ceiling, |current| current.min(ceiling)),
233            ),
234            purpose_bytes: Arc::clone(&self.purpose_bytes),
235        }
236    }
237
238    /// Clone this boundary while applying a maximum duration from its original start.
239    #[must_use]
240    pub fn with_timeout_ceiling(&self, timeout: Duration) -> Self {
241        let ceiling = self
242            .started_at
243            .checked_add(timeout)
244            .unwrap_or(self.started_at);
245        Self {
246            cancellation: self.cancellation.clone(),
247            started_at: self.started_at,
248            deadline: Some(
249                self.deadline
250                    .map_or(ceiling, |deadline| deadline.min(ceiling)),
251            ),
252            worker_ceiling: self.worker_ceiling,
253            purpose_bytes: Arc::clone(&self.purpose_bytes),
254        }
255    }
256
257    /// Consume authored-purpose bytes from this operation's aggregate budget.
258    ///
259    /// # Errors
260    ///
261    /// Returns [`IndexWorkFailure::ResourceLimitExceeded`] when this read and
262    /// every prior read through a clone would exceed `limit`.
263    pub fn consume_purpose_bytes(&self, limit: u64, bytes: u64) -> Result<u64, IndexWorkFailure> {
264        let updated =
265            self.purpose_bytes
266                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
267                    (bytes <= limit.saturating_sub(current)).then(|| current.saturating_add(bytes))
268                });
269        match updated {
270            Ok(previous) => Ok(previous.saturating_add(bytes)),
271            Err(previous) => Err(IndexWorkFailure::resource_limit(
272                IndexWorkStage::Publication,
273                IndexWorkResource::PurposeBytes,
274                limit,
275                previous.saturating_add(bytes),
276            )),
277        }
278    }
279
280    /// Request cooperative cancellation for every clone of this control.
281    pub fn cancel(&self) {
282        self.cancellation.cancel();
283    }
284
285    /// Borrow the exact cancellation signal shared by every operation worker.
286    #[must_use]
287    pub const fn cancellation(&self) -> &IndexCancellation {
288        &self.cancellation
289    }
290
291    /// Check cancellation and deadline state at one typed work stage.
292    ///
293    /// # Errors
294    ///
295    /// Returns [`IndexWorkFailure::Cancelled`] when cancellation was requested,
296    /// or [`IndexWorkFailure::DeadlineExceeded`] when the deadline elapsed.
297    pub fn check(&self, stage: IndexWorkStage) -> Result<(), IndexWorkFailure> {
298        if self.cancellation.is_cancelled() {
299            return Err(IndexWorkFailure::Cancelled { stage });
300        }
301        if self
302            .deadline
303            .is_some_and(|deadline| Instant::now() >= deadline)
304        {
305            return Err(IndexWorkFailure::DeadlineExceeded { stage });
306        }
307        Ok(())
308    }
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314
315    /// Clones must observe one cancellation flag, deadline, and resource boundary.
316    #[test]
317    fn work_control_shares_cancellation_and_deadline() {
318        let cancellation = IndexCancellation::new();
319        let control = IndexWorkControl::new(cancellation.clone(), None);
320        let worker = control.clone();
321        assert_eq!(worker.started_at(), control.started_at());
322        assert_eq!(worker.worker_ceiling(), None);
323        assert!(!control.cancellation().is_cancelled());
324
325        let worker_bounded = worker.with_worker_ceiling(4);
326        let tighter_worker_bound = worker_bounded.with_worker_ceiling(2);
327        assert_eq!(worker_bounded.worker_ceiling(), Some(4));
328        assert_eq!(tighter_worker_bound.worker_ceiling(), Some(2));
329        assert_eq!(
330            tighter_worker_bound
331                .with_timeout_ceiling(Duration::from_secs(1))
332                .worker_ceiling(),
333            Some(2)
334        );
335        cancellation.cancel();
336        assert!(control.cancellation().is_cancelled());
337        assert_eq!(
338            worker.check(IndexWorkStage::SourceHash),
339            Err(IndexWorkFailure::Cancelled {
340                stage: IndexWorkStage::SourceHash,
341            })
342        );
343
344        let bounded = tighter_worker_bound.with_timeout_ceiling(Duration::from_secs(1));
345        assert_eq!(bounded.started_at(), worker.started_at());
346        assert!(bounded.deadline().is_some());
347        assert_eq!(control.consume_purpose_bytes(8, 3), Ok(3));
348        assert_eq!(bounded.consume_purpose_bytes(8, 5), Ok(8));
349        assert_eq!(
350            worker.consume_purpose_bytes(8, 1),
351            Err(IndexWorkFailure::ResourceLimitExceeded {
352                stage: IndexWorkStage::Publication,
353                resource: IndexWorkResource::PurposeBytes,
354                limit: 8,
355                observed: 9,
356            })
357        );
358
359        let elapsed = IndexWorkControl::with_deadline(IndexCancellation::new(), Instant::now());
360        assert_eq!(
361            elapsed.check(IndexWorkStage::RepositoryTraversal),
362            Err(IndexWorkFailure::DeadlineExceeded {
363                stage: IndexWorkStage::RepositoryTraversal,
364            })
365        );
366
367        let overflow = IndexWorkControl::new(IndexCancellation::new(), Some(Duration::MAX));
368        assert_eq!(
369            overflow.check(IndexWorkStage::RepositoryTraversal),
370            Err(IndexWorkFailure::DeadlineExceeded {
371                stage: IndexWorkStage::RepositoryTraversal,
372            })
373        );
374    }
375}