mongreldb_core/
execution.rs

1//! Shared cooperative cancellation and deadline control.
2
3use crate::{MongrelError, Result};
4use std::sync::atomic::{AtomicU8, Ordering};
5use std::sync::Arc;
6use std::time::{Duration, Instant};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9#[repr(u8)]
10pub enum CancellationReason {
11    None = 0,
12    ClientRequest = 1,
13    Deadline = 2,
14    ClientDisconnected = 3,
15    SessionClosed = 4,
16    ServerShutdown = 5,
17}
18
19impl CancellationReason {
20    fn from_u8(value: u8) -> Self {
21        match value {
22            1 => Self::ClientRequest,
23            2 => Self::Deadline,
24            3 => Self::ClientDisconnected,
25            4 => Self::SessionClosed,
26            5 => Self::ServerShutdown,
27            _ => Self::None,
28        }
29    }
30}
31
32#[derive(Debug)]
33struct CancellationState {
34    reason: AtomicU8,
35}
36
37/// Cloneable cooperative control shared through one execution.
38///
39/// Child controls inherit every ancestor cancellation state and the tightest
40/// deadline. Cancelling a child does not cancel its parent.
41#[derive(Debug, Clone)]
42pub struct ExecutionControl {
43    states: Arc<Vec<Arc<CancellationState>>>,
44    own: Arc<CancellationState>,
45    wake: Arc<tokio::sync::Notify>,
46    deadline: Option<Instant>,
47}
48
49impl ExecutionControl {
50    pub fn new(deadline: Option<Instant>) -> Self {
51        let own = Arc::new(CancellationState {
52            reason: AtomicU8::new(CancellationReason::None as u8),
53        });
54        Self {
55            states: Arc::new(vec![Arc::clone(&own)]),
56            own,
57            wake: Arc::new(tokio::sync::Notify::new()),
58            deadline,
59        }
60    }
61
62    pub fn with_timeout(timeout: Duration) -> Self {
63        Self::new(Some(Instant::now() + timeout))
64    }
65
66    pub fn child_with_deadline(&self, deadline: Option<Instant>) -> Self {
67        let deadline = match (self.deadline, deadline) {
68            (Some(parent), Some(child)) => Some(parent.min(child)),
69            (Some(parent), None) => Some(parent),
70            (None, child) => child,
71        };
72        let own = Arc::new(CancellationState {
73            reason: AtomicU8::new(CancellationReason::None as u8),
74        });
75        let mut states = self.states.as_ref().clone();
76        states.push(Arc::clone(&own));
77        Self {
78            states: Arc::new(states),
79            own,
80            wake: Arc::clone(&self.wake),
81            deadline,
82        }
83    }
84
85    pub fn child_with_timeout(&self, timeout: Duration) -> Self {
86        self.child_with_deadline(Some(Instant::now() + timeout))
87    }
88
89    pub fn cancel(&self, reason: CancellationReason) {
90        if reason == CancellationReason::None {
91            return;
92        }
93        let _ = self.own.reason.compare_exchange(
94            CancellationReason::None as u8,
95            reason as u8,
96            Ordering::AcqRel,
97            Ordering::Acquire,
98        );
99        self.wake.notify_waiters();
100    }
101
102    pub fn checkpoint(&self) -> Result<()> {
103        if self
104            .deadline
105            .is_some_and(|deadline| Instant::now() >= deadline)
106        {
107            self.cancel(CancellationReason::Deadline);
108        }
109        match self.reason() {
110            CancellationReason::None => Ok(()),
111            CancellationReason::Deadline => Err(MongrelError::DeadlineExceeded),
112            _ => Err(MongrelError::Cancelled),
113        }
114    }
115
116    pub async fn cancelled(&self) {
117        loop {
118            if self.checkpoint().is_err() {
119                return;
120            }
121            let notified = self.wake.notified();
122            if self.checkpoint().is_err() {
123                return;
124            }
125            if let Some(remaining) = self.remaining_duration() {
126                if tokio::time::timeout(remaining, notified).await.is_err() {
127                    self.cancel(CancellationReason::Deadline);
128                    return;
129                }
130            } else {
131                notified.await;
132            }
133        }
134    }
135
136    pub fn remaining_duration(&self) -> Option<Duration> {
137        self.deadline
138            .map(|deadline| deadline.saturating_duration_since(Instant::now()))
139    }
140
141    pub fn deadline(&self) -> Option<Instant> {
142        self.deadline
143    }
144
145    pub fn reason(&self) -> CancellationReason {
146        self.states
147            .iter()
148            .map(|state| CancellationReason::from_u8(state.reason.load(Ordering::Acquire)))
149            .find(|reason| *reason != CancellationReason::None)
150            .unwrap_or(CancellationReason::None)
151    }
152
153    pub fn is_cancelled(&self) -> bool {
154        if self
155            .deadline
156            .is_some_and(|deadline| Instant::now() >= deadline)
157        {
158            self.cancel(CancellationReason::Deadline);
159        }
160        self.reason() != CancellationReason::None
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167
168    #[test]
169    fn first_reason_wins_and_unknown_values_are_checked() {
170        let control = ExecutionControl::new(None);
171        control.cancel(CancellationReason::ClientRequest);
172        control.cancel(CancellationReason::ServerShutdown);
173        assert_eq!(control.reason(), CancellationReason::ClientRequest);
174        assert!(matches!(control.checkpoint(), Err(MongrelError::Cancelled)));
175        assert_eq!(CancellationReason::from_u8(255), CancellationReason::None);
176    }
177
178    #[test]
179    fn deadline_returns_structured_error() {
180        let control = ExecutionControl::new(Some(Instant::now()));
181        assert!(matches!(
182            control.checkpoint(),
183            Err(MongrelError::DeadlineExceeded)
184        ));
185        assert_eq!(control.reason(), CancellationReason::Deadline);
186    }
187
188    #[test]
189    fn children_inherit_parent_but_do_not_cancel_it() {
190        let parent = ExecutionControl::new(None);
191        let child = parent.child_with_deadline(None);
192        child.cancel(CancellationReason::ClientRequest);
193        assert!(parent.checkpoint().is_ok());
194        assert!(child.checkpoint().is_err());
195
196        let sibling = parent.child_with_deadline(None);
197        parent.cancel(CancellationReason::SessionClosed);
198        assert_eq!(sibling.reason(), CancellationReason::SessionClosed);
199        assert!(sibling.checkpoint().is_err());
200    }
201
202    #[test]
203    fn child_deadline_cannot_weaken_parent() {
204        let parent_deadline = Instant::now() + Duration::from_secs(1);
205        let parent = ExecutionControl::new(Some(parent_deadline));
206        let child = parent.child_with_timeout(Duration::from_secs(60));
207        assert_eq!(child.deadline(), Some(parent_deadline));
208    }
209}
mongreldb_core/execution.rs

mongreldb_core/
execution.rs
