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    /// Return [`QueryCancelled`] if cancellation was signalled.
84    /// `Ok(())` otherwise.
85    ///
86    /// Designed for the inner loop of every operator: a single
87    /// relaxed-ordered atomic load on the happy path.
88    pub fn check(&self) -> Result<(), QueryCancelled> {
89        if self.is_cancelled() {
90            Err(QueryCancelled)
91        } else {
92            Ok(())
93        }
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use std::thread;
101
102    #[test]
103    fn fresh_token_is_not_cancelled() {
104        let tok = CancellationToken::new();
105        assert!(!tok.is_cancelled());
106        assert!(tok.check().is_ok());
107    }
108
109    #[test]
110    fn cancel_propagates_through_clone() {
111        let tok = CancellationToken::new();
112        let observer = tok.clone();
113        tok.cancel();
114        assert!(observer.is_cancelled());
115        assert_eq!(observer.check(), Err(QueryCancelled));
116    }
117
118    #[test]
119    fn reset_clears_signal() {
120        let tok = CancellationToken::new();
121        tok.cancel();
122        tok.reset();
123        assert!(!tok.is_cancelled());
124        assert!(tok.check().is_ok());
125    }
126
127    #[test]
128    fn cancel_visible_across_threads() {
129        let tok = CancellationToken::new();
130        let worker = tok.clone();
131        let handle = thread::spawn(move || {
132            // Spin until the parent cancels (test-only; real
133            // operators check at chunk boundaries instead).
134            while !worker.is_cancelled() {
135                std::hint::spin_loop();
136            }
137            worker.check()
138        });
139        tok.cancel();
140        let res = handle.join().unwrap();
141        assert_eq!(res, Err(QueryCancelled));
142    }
143}
uqa_core/cancel.rs

uqa_core/
cancel.rs
