1use std::fmt;
4use std::sync::{
5 Arc,
6 atomic::{AtomicBool, AtomicU64, Ordering},
7};
8use std::time::{Duration, Instant};
9use thiserror::Error;
10
11#[derive(Clone, Debug, Default)]
13pub struct IndexCancellation {
14 cancelled: Arc<AtomicBool>,
16}
17
18impl IndexCancellation {
19 #[must_use]
21 pub fn new() -> Self {
22 Self::default()
23 }
24
25 pub fn cancel(&self) {
27 self.cancelled.store(true, Ordering::Relaxed);
28 }
29
30 #[must_use]
32 pub fn is_cancelled(&self) -> bool {
33 self.cancelled.load(Ordering::Relaxed)
34 }
35}
36
37#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub enum IndexWorkStage {
40 RepositoryTraversal,
42 SourceMetadata,
44 SourceHash,
46 TextIndex,
48 SymbolParsing,
50 ScanFinalization,
52 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#[derive(Clone, Copy, Debug, Eq, PartialEq)]
72pub enum IndexWorkResource {
73 Entries,
75 SourceBytes,
77 TextBytes,
79 PurposeBytes,
81 PurposeRecords,
83 SymbolJobs,
85 SymbolRows,
87 RelationRows,
89 OutputBytes,
91 Workers,
93 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#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
117pub enum IndexWorkFailure {
118 #[error("index work was canceled during {stage}")]
120 Cancelled {
121 stage: IndexWorkStage,
123 },
124 #[error("index work deadline was reached during {stage}")]
126 DeadlineExceeded {
127 stage: IndexWorkStage,
129 },
130 #[error(
132 "index work exceeded the {resource} limit during {stage}: observed {observed}, limit {limit}"
133 )]
134 ResourceLimitExceeded {
135 stage: IndexWorkStage,
137 resource: IndexWorkResource,
139 limit: u64,
141 observed: u64,
143 },
144}
145
146impl IndexWorkFailure {
147 #[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#[derive(Clone, Debug)]
166pub struct IndexWorkControl {
167 cancellation: IndexCancellation,
169 started_at: Instant,
171 deadline: Option<Instant>,
173 worker_ceiling: Option<usize>,
175 purpose_bytes: Arc<AtomicU64>,
177}
178
179impl IndexWorkControl {
180 #[must_use]
182 pub fn new(cancellation: IndexCancellation, timeout: Option<Duration>) -> Self {
183 let started_at = Instant::now();
184 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 #[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 #[must_use]
209 pub fn started_at(&self) -> Instant {
210 self.started_at
211 }
212
213 #[must_use]
215 pub fn deadline(&self) -> Option<Instant> {
216 self.deadline
217 }
218
219 #[must_use]
221 pub fn worker_ceiling(&self) -> Option<usize> {
222 self.worker_ceiling
223 }
224
225 #[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 #[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 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 pub fn cancel(&self) {
285 self.cancellation.cancel();
286 }
287
288 #[must_use]
290 pub const fn cancellation(&self) -> &IndexCancellation {
291 &self.cancellation
292 }
293
294 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 #[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}