Skip to main content

diskann_quantization/
cancel.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! Utilities to support cancelation of long-running operations.
7
8use std::sync::atomic::{AtomicBool, Ordering};
9
10/// Provides a means for cancelling long-running operations.
11///
12/// Reasonable implementation of this trait should ensure that once `should_cancel` returns
13/// `true`, all future calls to `should_cancel` should **also** return `true`.
14pub trait Cancelation {
15    fn should_cancel(&self) -> bool;
16}
17
18/// A light-weight cancelation token based on an `AtomicBool`.
19pub struct AtomicCancelation<'a>(&'a AtomicBool);
20
21impl<'a> AtomicCancelation<'a> {
22    pub fn new(val: &'a AtomicBool) -> Self {
23        Self(val)
24    }
25}
26
27impl Cancelation for AtomicCancelation<'_> {
28    fn should_cancel(&self) -> bool {
29        self.0.load(Ordering::Relaxed)
30    }
31}
32
33/// A no-op cancelation token.
34pub struct DontCancel;
35
36impl Cancelation for DontCancel {
37    fn should_cancel(&self) -> bool {
38        false
39    }
40}
41
42///////////
43// Tests //
44///////////
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    #[test]
51    fn test_atomic_cancelation() {
52        let x = AtomicBool::new(false);
53        let y = AtomicCancelation::new(&x);
54        assert!(!y.should_cancel());
55        x.store(true, Ordering::Relaxed);
56        assert!(y.should_cancel());
57    }
58
59    #[test]
60    fn test_no_cancelation() {
61        let x = DontCancel;
62        assert!(!x.should_cancel());
63    }
64}