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}
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#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
114pub enum IndexWorkFailure {
115 #[error("index work was canceled during {stage}")]
117 Cancelled {
118 stage: IndexWorkStage,
120 },
121 #[error("index work deadline was reached during {stage}")]
123 DeadlineExceeded {
124 stage: IndexWorkStage,
126 },
127 #[error(
129 "index work exceeded the {resource} limit during {stage}: observed {observed}, limit {limit}"
130 )]
131 ResourceLimitExceeded {
132 stage: IndexWorkStage,
134 resource: IndexWorkResource,
136 limit: u64,
138 observed: u64,
140 },
141}
142
143impl IndexWorkFailure {
144 #[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#[derive(Clone, Debug)]
163pub struct IndexWorkControl {
164 cancellation: IndexCancellation,
166 started_at: Instant,
168 deadline: Option<Instant>,
170 worker_ceiling: Option<usize>,
172 purpose_bytes: Arc<AtomicU64>,
174}
175
176impl IndexWorkControl {
177 #[must_use]
179 pub fn new(cancellation: IndexCancellation, timeout: Option<Duration>) -> Self {
180 let started_at = Instant::now();
181 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 #[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 #[must_use]
206 pub fn started_at(&self) -> Instant {
207 self.started_at
208 }
209
210 #[must_use]
212 pub fn deadline(&self) -> Option<Instant> {
213 self.deadline
214 }
215
216 #[must_use]
218 pub fn worker_ceiling(&self) -> Option<usize> {
219 self.worker_ceiling
220 }
221
222 #[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 #[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 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 pub fn cancel(&self) {
282 self.cancellation.cancel();
283 }
284
285 #[must_use]
287 pub const fn cancellation(&self) -> &IndexCancellation {
288 &self.cancellation
289 }
290
291 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 #[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}