Skip to main content

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    /// Interpreter instructions admitted for one contained parser execution.
94    ParserFuel,
95}
96
97impl fmt::Display for IndexWorkResource {
98    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
99        formatter.write_str(match self {
100            Self::Entries => "entries",
101            Self::SourceBytes => "source_bytes",
102            Self::TextBytes => "text_bytes",
103            Self::PurposeBytes => "purpose_bytes",
104            Self::PurposeRecords => "purpose_records",
105            Self::SymbolJobs => "symbol_jobs",
106            Self::SymbolRows => "symbol_rows",
107            Self::RelationRows => "relation_rows",
108            Self::OutputBytes => "output_bytes",
109            Self::Workers => "workers",
110            Self::ParserFuel => "parser_fuel",
111        })
112    }
113}
114
115/// Typed cooperative-stop failures returned by bounded indexing work.
116#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
117pub enum IndexWorkFailure {
118    /// The caller requested cancellation.
119    #[error("index work was canceled during {stage}")]
120    Cancelled {
121        /// Stage observing cancellation.
122        stage: IndexWorkStage,
123    },
124    /// The operation reached its shared deadline.
125    #[error("index work deadline was reached during {stage}")]
126    DeadlineExceeded {
127        /// Stage observing the elapsed deadline.
128        stage: IndexWorkStage,
129    },
130    /// Work would exceed one declared resource limit.
131    #[error(
132        "index work exceeded the {resource} limit during {stage}: observed {observed}, limit {limit}"
133    )]
134    ResourceLimitExceeded {
135        /// Stage observing the exhausted resource.
136        stage: IndexWorkStage,
137        /// Resource whose limit was reached.
138        resource: IndexWorkResource,
139        /// Configured inclusive maximum.
140        limit: u64,
141        /// First observed value beyond the maximum.
142        observed: u64,
143    },
144}
145
146impl IndexWorkFailure {
147    /// Construct a typed resource-limit failure.
148    #[must_use]
149    pub const fn resource_limit(
150        stage: IndexWorkStage,
151        resource: IndexWorkResource,
152        limit: u64,
153        observed: u64,
154    ) -> Self {
155        Self::ResourceLimitExceeded {
156            stage,
157            resource,
158            limit,
159            observed,
160        }
161    }
162}
163
164/// Shared cancellation and deadline boundary for one indexing operation.
165#[derive(Clone, Debug)]
166pub struct IndexWorkControl {
167    /// Cooperative cancellation signal shared across operation workers.
168    cancellation: IndexCancellation,
169    /// Common operation start used to derive and report its deadline.
170    started_at: Instant,
171    /// Optional absolute deadline observed by every operation worker.
172    deadline: Option<Instant>,
173    /// Optional maximum worker count shared by every operation stage.
174    worker_ceiling: Option<usize>,
175    /// Authored-purpose bytes consumed by this operation and every clone.
176    purpose_bytes: Arc<AtomicU64>,
177}
178
179impl IndexWorkControl {
180    /// Create a work boundary with an optional timeout from one shared start.
181    #[must_use]
182    pub fn new(cancellation: IndexCancellation, timeout: Option<Duration>) -> Self {
183        let started_at = Instant::now();
184        // An unrepresentable deadline fails closed instead of silently removing the bound.
185        let deadline = timeout.map(|timeout| started_at.checked_add(timeout).unwrap_or(started_at));
186        Self {
187            cancellation,
188            started_at,
189            deadline,
190            worker_ceiling: None,
191            purpose_bytes: Arc::new(AtomicU64::new(0)),
192        }
193    }
194
195    /// Create a work boundary with a caller-selected absolute deadline.
196    #[must_use]
197    pub fn with_deadline(cancellation: IndexCancellation, deadline: Instant) -> Self {
198        Self {
199            cancellation,
200            started_at: Instant::now(),
201            deadline: Some(deadline),
202            worker_ceiling: None,
203            purpose_bytes: Arc::new(AtomicU64::new(0)),
204        }
205    }
206
207    /// Return the common start instant for this operation.
208    #[must_use]
209    pub fn started_at(&self) -> Instant {
210        self.started_at
211    }
212
213    /// Return the common absolute deadline, when configured.
214    #[must_use]
215    pub fn deadline(&self) -> Option<Instant> {
216        self.deadline
217    }
218
219    /// Return the maximum workers available to each operation stage, when bounded.
220    #[must_use]
221    pub fn worker_ceiling(&self) -> Option<usize> {
222        self.worker_ceiling
223    }
224
225    /// Clone this boundary while applying a worker ceiling to every operation stage.
226    #[must_use]
227    pub fn with_worker_ceiling(&self, max_workers: usize) -> Self {
228        let ceiling = max_workers.max(1);
229        Self {
230            cancellation: self.cancellation.clone(),
231            started_at: self.started_at,
232            deadline: self.deadline,
233            worker_ceiling: Some(
234                self.worker_ceiling
235                    .map_or(ceiling, |current| current.min(ceiling)),
236            ),
237            purpose_bytes: Arc::clone(&self.purpose_bytes),
238        }
239    }
240
241    /// Clone this boundary while applying a maximum duration from its original start.
242    #[must_use]
243    pub fn with_timeout_ceiling(&self, timeout: Duration) -> Self {
244        let ceiling = self
245            .started_at
246            .checked_add(timeout)
247            .unwrap_or(self.started_at);
248        Self {
249            cancellation: self.cancellation.clone(),
250            started_at: self.started_at,
251            deadline: Some(
252                self.deadline
253                    .map_or(ceiling, |deadline| deadline.min(ceiling)),
254            ),
255            worker_ceiling: self.worker_ceiling,
256            purpose_bytes: Arc::clone(&self.purpose_bytes),
257        }
258    }
259
260    /// Consume authored-purpose bytes from this operation's aggregate budget.
261    ///
262    /// # Errors
263    ///
264    /// Returns [`IndexWorkFailure::ResourceLimitExceeded`] when this read and
265    /// every prior read through a clone would exceed `limit`.
266    pub fn consume_purpose_bytes(&self, limit: u64, bytes: u64) -> Result<u64, IndexWorkFailure> {
267        let updated =
268            self.purpose_bytes
269                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
270                    (bytes <= limit.saturating_sub(current)).then(|| current.saturating_add(bytes))
271                });
272        match updated {
273            Ok(previous) => Ok(previous.saturating_add(bytes)),
274            Err(previous) => Err(IndexWorkFailure::resource_limit(
275                IndexWorkStage::Publication,
276                IndexWorkResource::PurposeBytes,
277                limit,
278                previous.saturating_add(bytes),
279            )),
280        }
281    }
282
283    /// Request cooperative cancellation for every clone of this control.
284    pub fn cancel(&self) {
285        self.cancellation.cancel();
286    }
287
288    /// Borrow the exact cancellation signal shared by every operation worker.
289    #[must_use]
290    pub const fn cancellation(&self) -> &IndexCancellation {
291        &self.cancellation
292    }
293
294    /// Check cancellation and deadline state at one typed work stage.
295    ///
296    /// # Errors
297    ///
298    /// Returns [`IndexWorkFailure::Cancelled`] when cancellation was requested,
299    /// or [`IndexWorkFailure::DeadlineExceeded`] when the deadline elapsed.
300    pub fn check(&self, stage: IndexWorkStage) -> Result<(), IndexWorkFailure> {
301        if self.cancellation.is_cancelled() {
302            return Err(IndexWorkFailure::Cancelled { stage });
303        }
304        if self
305            .deadline
306            .is_some_and(|deadline| Instant::now() >= deadline)
307        {
308            return Err(IndexWorkFailure::DeadlineExceeded { stage });
309        }
310        Ok(())
311    }
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317
318    /// Clones must observe one cancellation flag, deadline, and resource boundary.
319    #[test]
320    fn work_control_shares_cancellation_and_deadline() {
321        let cancellation = IndexCancellation::new();
322        let control = IndexWorkControl::new(cancellation.clone(), None);
323        let worker = control.clone();
324        assert_eq!(worker.started_at(), control.started_at());
325        assert_eq!(worker.worker_ceiling(), None);
326        assert!(!control.cancellation().is_cancelled());
327
328        let worker_bounded = worker.with_worker_ceiling(4);
329        let tighter_worker_bound = worker_bounded.with_worker_ceiling(2);
330        assert_eq!(worker_bounded.worker_ceiling(), Some(4));
331        assert_eq!(tighter_worker_bound.worker_ceiling(), Some(2));
332        assert_eq!(
333            tighter_worker_bound
334                .with_timeout_ceiling(Duration::from_secs(1))
335                .worker_ceiling(),
336            Some(2)
337        );
338        cancellation.cancel();
339        assert!(control.cancellation().is_cancelled());
340        assert_eq!(
341            worker.check(IndexWorkStage::SourceHash),
342            Err(IndexWorkFailure::Cancelled {
343                stage: IndexWorkStage::SourceHash,
344            })
345        );
346
347        let bounded = tighter_worker_bound.with_timeout_ceiling(Duration::from_secs(1));
348        assert_eq!(bounded.started_at(), worker.started_at());
349        assert!(bounded.deadline().is_some());
350        assert_eq!(control.consume_purpose_bytes(8, 3), Ok(3));
351        assert_eq!(bounded.consume_purpose_bytes(8, 5), Ok(8));
352        assert_eq!(
353            worker.consume_purpose_bytes(8, 1),
354            Err(IndexWorkFailure::ResourceLimitExceeded {
355                stage: IndexWorkStage::Publication,
356                resource: IndexWorkResource::PurposeBytes,
357                limit: 8,
358                observed: 9,
359            })
360        );
361
362        let elapsed = IndexWorkControl::with_deadline(IndexCancellation::new(), Instant::now());
363        assert_eq!(
364            elapsed.check(IndexWorkStage::RepositoryTraversal),
365            Err(IndexWorkFailure::DeadlineExceeded {
366                stage: IndexWorkStage::RepositoryTraversal,
367            })
368        );
369
370        let overflow = IndexWorkControl::new(IndexCancellation::new(), Some(Duration::MAX));
371        assert_eq!(
372            overflow.check(IndexWorkStage::RepositoryTraversal),
373            Err(IndexWorkFailure::DeadlineExceeded {
374                stage: IndexWorkStage::RepositoryTraversal,
375            })
376        );
377    }
378}