uqa_core/
cancel.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Query cancellation support.
8//!
9//! A [`CancellationToken`] is a cheap-to-clone, thread-safe one-shot
10//! flag stored on `Engine` and propagated into every
11//! `PhysicalOperator` / `Operator` hot loop. Operators call
12//! [`CancellationToken::check`] at chunk boundaries; if the flag has
13//! been set from another thread, `check` returns
14//! [`QueryCancelled`] which surfaces to the SQL layer as
15//! `PostgreSQL` `SQLSTATE 57014` (`query_canceled`).
16//!
17//! ```rust
18//! use uqa_core::cancel::{CancellationToken, QueryCancelled};
19//!
20//! let tok = CancellationToken::new();
21//! let probe = tok.clone();
22//! tok.cancel();
23//! assert!(probe.is_cancelled());
24//! assert!(matches!(probe.check(), Err(QueryCancelled)));
25//! ```
26//!
27//! The token is a `Clone`-by-`Arc` handle: every clone speaks to the
28//! same underlying flag, so issuing `engine.cancel()` from one thread
29//! is immediately visible to any operator that received a clone of
30//! the token before the cancellation.
31
32use std::sync::atomic::{AtomicBool, Ordering};
33use std::sync::Arc;
34
35use thiserror::Error;
36
37/// Raised when a query is cancelled by user request. Matches
38/// `PostgreSQL` `SQLSTATE 57014` (`query_canceled`); its `Display` payload
39/// stays stable for log processing.
40#[derive(Debug, Clone, Copy, Error, PartialEq, Eq)]
41#[error("canceling statement due to user request")]
42pub struct QueryCancelled;
43
44/// `PostgreSQL` SQLSTATE for [`QueryCancelled`].
45pub const SQLSTATE_QUERY_CANCELED: &str = "57014";
46
47/// Thread-safe cancellation token for query execution.
48///
49/// Uses an [`AtomicBool`] behind an [`Arc`] so cloning is `O(1)` and
50/// every clone observes the same cancellation flag. Once
51/// [`CancellationToken::cancel`] has been called, every subsequent
52/// [`CancellationToken::check`] returns [`QueryCancelled`] until
53/// [`CancellationToken::reset`] is called.
54#[derive(Debug, Clone, Default)]
55pub struct CancellationToken {
56    flag: Arc<AtomicBool>,
57}
58
59impl CancellationToken {
60    pub fn new() -> Self {
61        Self {
62            flag: Arc::new(AtomicBool::new(false)),
63        }
64    }
65
66    /// Signal cancellation. Subsequent [`Self::check`] / [`Self::is_cancelled`]
67    /// observe the flag as set across all clones of this token.
68    pub fn cancel(&self) {
69        self.flag.store(true, Ordering::Release);
70    }
71
72    /// Clear the cancellation signal for the next query. Operators
73    /// holding a clone of this token through their lifetime see the
74    /// reset on the next `check`.
75    pub fn reset(&self) {
76        self.flag.store(false, Ordering::Release);
77    }
78
79    pub fn is_cancelled(&self) -> bool {
80        self.flag.load(Ordering::Acquire)
81    }
82
83    /// Whether both handles observe the same signal, independently of its current cancelled state.
84    pub fn shares_signal(&self, other: &Self) -> bool {
85        Arc::ptr_eq(&self.flag, &other.flag)
86    }
87
88    /// Return [`QueryCancelled`] if cancellation was signalled.
89    /// `Ok(())` otherwise.
90    ///
91    /// Designed for the inner loop of every operator: a single
92    /// relaxed-ordered atomic load on the happy path.
93    pub fn check(&self) -> Result<(), QueryCancelled> {
94        if self.is_cancelled() {
95            Err(QueryCancelled)
96        } else {
97            Ok(())
98        }
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use std::thread;
106
107    #[test]
108    fn fresh_token_is_not_cancelled() {
109        let tok = CancellationToken::new();
110        assert!(!tok.is_cancelled());
111        assert!(tok.check().is_ok());
112    }
113
114    #[test]
115    fn cancel_propagates_through_clone() {
116        let tok = CancellationToken::new();
117        let observer = tok.clone();
118        tok.cancel();
119        assert!(observer.is_cancelled());
120        assert_eq!(observer.check(), Err(QueryCancelled));
121    }
122
123    #[test]
124    fn signal_identity_distinguishes_independent_tokens_with_equal_states() {
125        let first = CancellationToken::new();
126        let retained = first.clone();
127        let independent = CancellationToken::new();
128        assert!(first.shares_signal(&retained));
129        assert!(!first.shares_signal(&independent));
130        first.cancel();
131        independent.cancel();
132        assert!(first.shares_signal(&retained));
133        assert!(!first.shares_signal(&independent));
134        first.reset();
135        assert!(first.shares_signal(&retained));
136    }
137
138    #[test]
139    fn reset_clears_signal() {
140        let tok = CancellationToken::new();
141        tok.cancel();
142        tok.reset();
143        assert!(!tok.is_cancelled());
144        assert!(tok.check().is_ok());
145    }
146
147    #[test]
148    fn cancel_visible_across_threads() {
149        let tok = CancellationToken::new();
150        let worker = tok.clone();
151        let handle = thread::spawn(move || {
152            // Spin until the parent cancels (test-only; real
153            // operators check at chunk boundaries instead).
154            while !worker.is_cancelled() {
155                std::hint::spin_loop();
156            }
157            worker.check()
158        });
159        tok.cancel();
160        let res = handle.join().unwrap();
161        assert_eq!(res, Err(QueryCancelled));
162    }
163}
uqa_core/cancel.rs

uqa_core/
cancel.rs
