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 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#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
120pub enum IndexWorkFailure {
121 #[error("index work was canceled during {stage}")]
123 Cancelled {
124 stage: IndexWorkStage,
126 },
127 #[error("index work deadline was reached during {stage}")]
129 DeadlineExceeded {
130 stage: IndexWorkStage,
132 },
133 #[error(
135 "index work exceeded the {resource} limit during {stage}: observed {observed}, limit {limit}"
136 )]
137 ResourceLimitExceeded {
138 stage: IndexWorkStage,
140 resource: IndexWorkResource,
142 limit: u64,
144 observed: u64,
146 },
147}
148
149impl IndexWorkFailure {
150 #[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#[derive(Clone, Debug)]
169pub struct IndexWorkControl {
170 cancellation: IndexCancellation,
172 started_at: Instant,
174 deadline: Option<Instant>,
176 worker_ceiling: Option<usize>,
178 purpose_bytes: Arc<AtomicU64>,
180}
181
182impl IndexWorkControl {
183 #[must_use]
185 pub fn new(cancellation: IndexCancellation, timeout: Option<Duration>) -> Self {
186 let started_at = Instant::now();
187 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 #[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 #[must_use]
212 pub fn started_at(&self) -> Instant {
213 self.started_at
214 }
215
216 #[must_use]
218 pub fn deadline(&self) -> Option<Instant> {
219 self.deadline
220 }
221
222 #[must_use]
224 pub fn worker_ceiling(&self) -> Option<usize> {
225 self.worker_ceiling
226 }
227
228 #[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 #[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 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 pub fn cancel(&self) {
288 self.cancellation.cancel();
289 }
290
291 #[must_use]
293 pub const fn cancellation(&self) -> &IndexCancellation {
294 &self.cancellation
295 }
296
297 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 #[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}