1use super::{ensure_same_generation, ensure_writable, persist_index_cache, Collection};
4use crate::error::{Error, Result};
5use crate::index::IndexRegistry;
6use serde::{Deserialize, Serialize};
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::{Arc, Condvar, Mutex, MutexGuard};
9use std::thread::{self, JoinHandle};
10use std::time::Duration;
11
12const DEFAULT_INTERVAL: Duration = Duration::from_secs(60);
13const MIN_INTERVAL: Duration = Duration::from_millis(10);
14const MAX_INTERVAL: Duration = Duration::from_secs(365 * 24 * 60 * 60);
15const MAX_ERROR_BYTES: usize = 1_024;
16
17#[derive(Clone, Copy, Debug, PartialEq, Eq)]
19#[must_use = "maintenance options do nothing until passed to Collection::start_maintenance"]
20pub struct CollectionMaintenanceOptions {
21 interval: Duration,
22}
23
24impl CollectionMaintenanceOptions {
25 pub fn new() -> Self {
26 Self::default()
27 }
28
29 pub fn try_with_interval(mut self, interval: Duration) -> Result<Self> {
32 validate_interval(interval)?;
33 self.interval = interval;
34 Ok(self)
35 }
36
37 pub fn interval(&self) -> Duration {
38 self.interval
39 }
40}
41
42impl Default for CollectionMaintenanceOptions {
43 fn default() -> Self {
44 Self {
45 interval: DEFAULT_INTERVAL,
46 }
47 }
48}
49
50#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
52#[serde(rename_all = "snake_case")]
53pub enum CollectionMaintenancePhase {
54 Running,
55 Degraded,
56 Closing,
57 Closed,
58}
59
60#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
62pub struct CollectionMaintenanceHealth {
63 pub phase: CollectionMaintenancePhase,
64 pub interval_ms: u64,
65 pub worker_alive: bool,
66 pub run_in_progress: bool,
67 pub successful_runs: u64,
68 pub failed_runs: u64,
69 pub skipped_runs: u64,
70 pub last_attempted_revision: Option<u64>,
71 pub last_successful_revision: Option<u64>,
72 pub last_error: Option<String>,
73}
74
75impl CollectionMaintenanceHealth {
76 pub fn is_healthy(&self) -> bool {
78 self.last_error.is_none()
79 && matches!(
80 self.phase,
81 CollectionMaintenancePhase::Running | CollectionMaintenancePhase::Closed
82 )
83 }
84}
85
86#[derive(Debug)]
87struct MaintenanceState {
88 health: CollectionMaintenanceHealth,
89 stop_requested: bool,
90 run_requested: bool,
91}
92
93#[derive(Debug)]
94struct MaintenanceShared {
95 state: Mutex<MaintenanceState>,
96 wake: Condvar,
97}
98
99#[must_use = "retain and close the runtime to own its background worker"]
106pub struct CollectionMaintenanceRuntime {
107 collection: Collection,
108 shared: Arc<MaintenanceShared>,
109 worker: Mutex<Option<JoinHandle<()>>>,
110 close_gate: Mutex<()>,
111 claim_released: AtomicBool,
112}
113
114impl std::fmt::Debug for CollectionMaintenanceRuntime {
115 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
116 formatter
117 .debug_struct("CollectionMaintenanceRuntime")
118 .field("health", &self.health())
119 .finish_non_exhaustive()
120 }
121}
122
123impl CollectionMaintenanceRuntime {
124 pub fn health(&self) -> CollectionMaintenanceHealth {
126 lock_state(&self.shared).health.clone()
127 }
128
129 pub fn trigger(&self) -> Result<()> {
131 let mut state = lock_state(&self.shared);
132 if !state.health.worker_alive
133 || matches!(
134 state.health.phase,
135 CollectionMaintenancePhase::Closing | CollectionMaintenancePhase::Closed
136 )
137 {
138 return Err(Error::failed_precondition(
139 "background maintenance is not running",
140 ));
141 }
142 state.run_requested = true;
143 drop(state);
144 self.shared.wake.notify_one();
145 Ok(())
146 }
147
148 pub fn close(&self) -> Result<()> {
151 self.shutdown()
152 }
153
154 fn shutdown(&self) -> Result<()> {
155 let _close = lock_mutex(&self.close_gate);
156 {
157 let mut state = lock_state(&self.shared);
158 if state.health.phase == CollectionMaintenancePhase::Closed {
159 self.release_claim();
160 return Ok(());
161 }
162 state.stop_requested = true;
163 state.run_requested = false;
164 state.health.phase = CollectionMaintenancePhase::Closing;
165 }
166 self.shared.wake.notify_all();
167
168 let worker = lock_mutex(&self.worker).take();
169 let join_error = worker
170 .and_then(|worker| worker.join().err())
171 .map(|_| Error::internal("background maintenance worker panicked while shutting down"));
172 self.release_claim();
173
174 let mut state = lock_state(&self.shared);
175 state.health.worker_alive = false;
176 state.health.run_in_progress = false;
177 state.health.phase = CollectionMaintenancePhase::Closed;
178 if let Some(error) = &join_error {
179 state.health.failed_runs = state.health.failed_runs.saturating_add(1);
180 state.health.last_error = Some(bounded_error(error.to_string()));
181 }
182 drop(state);
183
184 if let Some(error) = join_error {
185 Err(error)
186 } else {
187 Ok(())
188 }
189 }
190
191 fn release_claim(&self) {
192 if !self.claim_released.swap(true, Ordering::AcqRel) {
193 self.collection
194 .inner
195 .maintenance_claimed
196 .store(false, Ordering::Release);
197 }
198 }
199}
200
201impl Drop for CollectionMaintenanceRuntime {
202 fn drop(&mut self) {
203 let _ = self.shutdown();
204 }
205}
206
207impl Collection {
208 pub fn start_maintenance(
214 &self,
215 options: CollectionMaintenanceOptions,
216 ) -> Result<CollectionMaintenanceRuntime> {
217 validate_interval(options.interval)?;
218 self.ensure_open()?;
219 {
220 let state = self
221 .inner
222 .state
223 .read()
224 .map_err(|_| Error::internal("collection state lock poisoned"))?;
225 ensure_writable(&state.options)?;
226 }
227 if self
228 .inner
229 .maintenance_claimed
230 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
231 .is_err()
232 {
233 return Err(Error::already_exists(
234 "background maintenance already has an owner",
235 ));
236 }
237 if !self.is_open() {
238 self.inner
239 .maintenance_claimed
240 .store(false, Ordering::Release);
241 return Err(Error::failed_precondition("collection is closed"));
242 }
243
244 let shared = Arc::new(MaintenanceShared {
245 state: Mutex::new(MaintenanceState {
246 health: CollectionMaintenanceHealth {
247 phase: CollectionMaintenancePhase::Running,
248 interval_ms: duration_ms(options.interval),
249 worker_alive: true,
250 run_in_progress: false,
251 successful_runs: 0,
252 failed_runs: 0,
253 skipped_runs: 0,
254 last_attempted_revision: None,
255 last_successful_revision: None,
256 last_error: None,
257 },
258 stop_requested: false,
259 run_requested: false,
260 }),
261 wake: Condvar::new(),
262 });
263 let worker_collection = self.clone();
264 let worker_shared = Arc::clone(&shared);
265 let worker = match thread::Builder::new()
266 .name("a3s-vec-maintenance".to_string())
267 .spawn(move || {
268 maintenance_worker(&worker_collection, &worker_shared, options.interval);
269 }) {
270 Ok(worker) => worker,
271 Err(error) => {
272 self.inner
273 .maintenance_claimed
274 .store(false, Ordering::Release);
275 return Err(Error::internal(format!(
276 "spawn background maintenance worker: {error}"
277 )));
278 }
279 };
280
281 Ok(CollectionMaintenanceRuntime {
282 collection: self.clone(),
283 shared,
284 worker: Mutex::new(Some(worker)),
285 close_gate: Mutex::new(()),
286 claim_released: AtomicBool::new(false),
287 })
288 }
289
290 fn run_maintenance_pass(&self) -> Result<u64> {
291 self.ensure_open()?;
292 let _writer = self
293 .inner
294 .writer
295 .lock()
296 .map_err(|_| Error::internal("writer lock poisoned"))?;
297 let current = self
298 .inner
299 .state
300 .read()
301 .map_err(|_| Error::internal("collection state lock poisoned"))?
302 .clone();
303 ensure_writable(¤t.options)?;
304
305 let indexes = Arc::new(IndexRegistry::build(
306 ¤t.schema,
307 ¤t.docs,
308 current.revision,
309 )?);
310 let resource_usage = match current.options.resource_limits.enforce_state(
311 ¤t.schema,
312 ¤t.docs,
313 &indexes,
314 ) {
315 Ok(usage) => usage,
316 Err(error) => {
317 current.stats.record_resource_limit_rejection();
318 return Err(error);
319 }
320 };
321 let mut state = self
322 .inner
323 .state
324 .write()
325 .map_err(|_| Error::internal("collection state lock poisoned"))?;
326 ensure_same_generation(&state, ¤t)?;
327 state.indexes = Arc::clone(&indexes);
328 state.resource_usage = resource_usage;
329 drop(state);
330
331 let mut storage = self
332 .inner
333 .storage
334 .lock()
335 .map_err(|_| Error::internal("storage lock poisoned"))?;
336 storage.checkpoint(
337 ¤t.schema,
338 current.docs.as_ref(),
339 current.revision,
340 true,
341 )?;
342 persist_index_cache(&storage, ¤t.schema, &indexes, current.revision, true);
343 Ok(current.revision)
344 }
345}
346
347fn maintenance_worker(
348 collection: &Collection,
349 shared: &Arc<MaintenanceShared>,
350 interval: Duration,
351) {
352 loop {
353 let state = lock_state(shared);
354 let waited = shared.wake.wait_timeout_while(state, interval, |state| {
355 !state.stop_requested && !state.run_requested
356 });
357 let (mut state, timeout) = match waited {
358 Ok(value) => value,
359 Err(poisoned) => poisoned.into_inner(),
360 };
361 if state.stop_requested {
362 break;
363 }
364 if !state.run_requested && !timeout.timed_out() {
365 continue;
366 }
367 state.run_requested = false;
368 drop(state);
369
370 let observed = collection.stats();
371 let mut state = lock_state(shared);
372 if state.stop_requested {
373 break;
374 }
375 let revision = match observed {
376 Ok(collection_stats) => collection_stats.revision,
377 Err(error) => {
378 record_failure(&mut state.health, None, &error);
379 let collection_closed = !collection.is_open();
380 drop(state);
381 if collection_closed {
382 break;
383 }
384 continue;
385 }
386 };
387 state.health.last_attempted_revision = Some(revision);
388 if state.health.last_successful_revision == Some(revision)
389 && state.health.last_error.is_none()
390 {
391 state.health.skipped_runs = state.health.skipped_runs.saturating_add(1);
392 continue;
393 }
394 state.health.run_in_progress = true;
395 drop(state);
396
397 let result = collection.run_maintenance_pass();
398 let mut state = lock_state(shared);
399 state.health.run_in_progress = false;
400 match result {
401 Ok(maintained_revision) => {
402 state.health.successful_runs = state.health.successful_runs.saturating_add(1);
403 state.health.last_successful_revision = Some(maintained_revision);
404 state.health.last_error = None;
405 if !state.stop_requested {
406 state.health.phase = CollectionMaintenancePhase::Running;
407 }
408 }
409 Err(error) => record_failure(&mut state.health, Some(revision), &error),
410 }
411 let collection_closed = !collection.is_open();
412 drop(state);
413 if collection_closed {
414 break;
415 }
416 }
417
418 let mut state = lock_state(shared);
419 state.health.worker_alive = false;
420 state.health.run_in_progress = false;
421 if !state.stop_requested {
422 state.health.phase = CollectionMaintenancePhase::Degraded;
423 if state.health.last_error.is_none() {
424 state.health.last_error = Some("background maintenance worker stopped".to_string());
425 }
426 }
427}
428
429fn record_failure(health: &mut CollectionMaintenanceHealth, revision: Option<u64>, error: &Error) {
430 health.failed_runs = health.failed_runs.saturating_add(1);
431 health.last_attempted_revision = revision;
432 health.last_error = Some(bounded_error(error.to_string()));
433 health.phase = CollectionMaintenancePhase::Degraded;
434}
435
436fn validate_interval(interval: Duration) -> Result<()> {
437 if interval < MIN_INTERVAL {
438 return Err(Error::invalid_argument(
439 "maintenance interval must be at least 10 milliseconds",
440 ));
441 }
442 if interval > MAX_INTERVAL {
443 return Err(Error::invalid_argument(
444 "maintenance interval must not exceed 365 days",
445 ));
446 }
447 Ok(())
448}
449
450fn duration_ms(duration: Duration) -> u64 {
451 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
452}
453
454fn bounded_error(mut message: String) -> String {
455 if message.len() <= MAX_ERROR_BYTES {
456 return message;
457 }
458 let mut end = MAX_ERROR_BYTES;
459 while !message.is_char_boundary(end) {
460 end -= 1;
461 }
462 message.truncate(end);
463 message
464}
465
466fn lock_state(shared: &MaintenanceShared) -> MutexGuard<'_, MaintenanceState> {
467 lock_mutex(&shared.state)
468}
469
470fn lock_mutex<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
471 match mutex.lock() {
472 Ok(guard) => guard,
473 Err(poisoned) => poisoned.into_inner(),
474 }
475}
476
477#[cfg(test)]
478mod tests {
479 use super::*;
480
481 #[test]
482 fn bounded_errors_preserve_utf8() {
483 let message = "é".repeat(MAX_ERROR_BYTES);
484 let bounded = bounded_error(message);
485 assert!(bounded.len() <= MAX_ERROR_BYTES);
486 assert!(bounded.chars().all(|character| character == 'é'));
487 }
488
489 #[test]
490 fn validate_interval_rejects_out_of_range_durations() {
491 assert!(validate_interval(Duration::from_millis(1)).is_err());
492 assert!(validate_interval(MIN_INTERVAL).is_ok());
493 assert!(validate_interval(MAX_INTERVAL).is_ok());
494 assert!(validate_interval(MAX_INTERVAL + Duration::from_secs(1)).is_err());
495 assert_eq!(duration_ms(Duration::from_millis(12)), 12);
496 }
497
498 #[test]
499 fn runtime_contract_is_send_and_sync() {
500 fn assert_send_sync<T: Send + Sync>() {}
501 assert_send_sync::<CollectionMaintenanceRuntime>();
502 }
503
504 #[test]
505 fn options_and_health_surface_interval_debug_and_closed_health() {
506 let options = CollectionMaintenanceOptions::new();
507 assert_eq!(options.interval(), DEFAULT_INTERVAL);
508 let tuned = options
509 .try_with_interval(Duration::from_millis(25))
510 .expect("interval");
511 assert_eq!(tuned.interval(), Duration::from_millis(25));
512
513 let healthy = CollectionMaintenanceHealth {
514 phase: CollectionMaintenancePhase::Running,
515 interval_ms: 25,
516 worker_alive: true,
517 run_in_progress: false,
518 successful_runs: 1,
519 failed_runs: 0,
520 skipped_runs: 0,
521 last_attempted_revision: Some(1),
522 last_successful_revision: Some(1),
523 last_error: None,
524 };
525 assert!(healthy.is_healthy());
526 let closed = CollectionMaintenanceHealth {
527 phase: CollectionMaintenancePhase::Closed,
528 ..healthy.clone()
529 };
530 assert!(closed.is_healthy());
531 let degraded = CollectionMaintenanceHealth {
532 phase: CollectionMaintenancePhase::Degraded,
533 last_error: Some("boom".into()),
534 ..healthy
535 };
536 assert!(!degraded.is_healthy());
537 let debug = format!("{degraded:?}");
538 assert!(debug.contains("Degraded") || debug.contains("boom"));
539 }
540}