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