hawdb_core/
cancellation.rs

1// Copyright 2026 Nowledge
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::BTreeMap;
16use std::error::Error;
17use std::fmt::Debug;
18use std::fmt::{self, Display, Formatter};
19use std::future::Future;
20use std::num::NonZeroUsize;
21use std::pin::Pin;
22use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
23use std::sync::{Arc, Mutex, Weak};
24use std::task::{Context, Poll, Waker};
25use std::time::{Duration, Instant};
26
27#[derive(Debug, Clone)]
28pub struct RuntimeCancellationToken {
29    state: Arc<RuntimeCancellationState>,
30}
31
32#[derive(Debug)]
33struct RuntimeCancellationState {
34    cancelled: AtomicBool,
35    parent: Option<RuntimeCancellationToken>,
36    next_waiter_id: AtomicU64,
37    waiters: Mutex<BTreeMap<u64, Waker>>,
38}
39
40impl Default for RuntimeCancellationToken {
41    fn default() -> Self {
42        Self {
43            state: Arc::new(RuntimeCancellationState {
44                cancelled: AtomicBool::new(false),
45                parent: None,
46                next_waiter_id: AtomicU64::new(1),
47                waiters: Mutex::new(BTreeMap::new()),
48            }),
49        }
50    }
51}
52
53impl RuntimeCancellationToken {
54    pub fn new() -> Self {
55        Self::default()
56    }
57
58    pub fn cancel(&self) -> bool {
59        if self.state.cancelled.swap(true, Ordering::AcqRel) {
60            return false;
61        }
62        let waiters = {
63            let mut waiters = self
64                .state
65                .waiters
66                .lock()
67                .unwrap_or_else(|poisoned| poisoned.into_inner());
68            std::mem::take(&mut *waiters)
69        };
70        for (_, waiter) in waiters {
71            waiter.wake();
72        }
73        true
74    }
75
76    pub fn is_cancelled(&self) -> bool {
77        self.state.cancelled.load(Ordering::Acquire)
78            || self
79                .state
80                .parent
81                .as_ref()
82                .is_some_and(RuntimeCancellationToken::is_cancelled)
83    }
84
85    pub fn child(&self) -> Self {
86        Self {
87            state: Arc::new(RuntimeCancellationState {
88                cancelled: AtomicBool::new(false),
89                parent: Some(self.clone()),
90                next_waiter_id: AtomicU64::new(1),
91                waiters: Mutex::new(BTreeMap::new()),
92            }),
93        }
94    }
95
96    /// Resolves when this token or any parent token is cancelled.
97    pub fn cancelled(&self) -> RuntimeCancellationFuture<'_> {
98        RuntimeCancellationFuture {
99            token: self,
100            registrations: Vec::new(),
101        }
102    }
103}
104
105struct RuntimeCancellationRegistration {
106    state: Weak<RuntimeCancellationState>,
107    waiter_id: u64,
108}
109
110/// A runtime-neutral future returned by [`RuntimeCancellationToken::cancelled`].
111pub struct RuntimeCancellationFuture<'a> {
112    token: &'a RuntimeCancellationToken,
113    registrations: Vec<RuntimeCancellationRegistration>,
114}
115
116impl Future for RuntimeCancellationFuture<'_> {
117    type Output = ();
118
119    fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
120        if self.token.is_cancelled() {
121            return Poll::Ready(());
122        }
123
124        if self.registrations.is_empty() {
125            let mut token = Some(self.token);
126            while let Some(current) = token {
127                let waiter_id = current.state.next_waiter_id.fetch_add(1, Ordering::Relaxed);
128                current
129                    .state
130                    .waiters
131                    .lock()
132                    .unwrap_or_else(|poisoned| poisoned.into_inner())
133                    .insert(waiter_id, context.waker().clone());
134                self.registrations.push(RuntimeCancellationRegistration {
135                    state: Arc::downgrade(&current.state),
136                    waiter_id,
137                });
138                token = current.state.parent.as_ref();
139            }
140        } else {
141            for registration in &self.registrations {
142                if let Some(state) = registration.state.upgrade()
143                    && let Some(waiter) = state
144                        .waiters
145                        .lock()
146                        .unwrap_or_else(|poisoned| poisoned.into_inner())
147                        .get_mut(&registration.waiter_id)
148                    && !waiter.will_wake(context.waker())
149                {
150                    *waiter = context.waker().clone();
151                }
152            }
153        }
154
155        if self.token.is_cancelled() {
156            Poll::Ready(())
157        } else {
158            Poll::Pending
159        }
160    }
161}
162
163impl Drop for RuntimeCancellationFuture<'_> {
164    fn drop(&mut self) {
165        for registration in &self.registrations {
166            if let Some(state) = registration.state.upgrade() {
167                state
168                    .waiters
169                    .lock()
170                    .unwrap_or_else(|poisoned| poisoned.into_inner())
171                    .remove(&registration.waiter_id);
172            }
173        }
174    }
175}
176
177/// A drop guard for one admitted storage I/O wave.
178pub trait RuntimeIoWavePermit: Debug + Send {}
179
180impl<T: Debug + Send> RuntimeIoWavePermit for T {}
181
182/// Execution-owned hook for acquiring the I/O capacity declared at admission.
183pub trait RuntimeIoWaveController: Debug + Send + Sync {
184    /// Attempts to acquire one I/O wave without blocking the calling thread.
185    ///
186    /// `Ok(None)` means that the reservation is currently saturated. Async
187    /// runtime adapters must yield before retrying instead of spinning.
188    fn try_acquire(
189        &self,
190        slots: NonZeroUsize,
191        context: &RuntimeTaskContext,
192    ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError>;
193
194    fn acquire(
195        &self,
196        slots: NonZeroUsize,
197        context: &RuntimeTaskContext,
198    ) -> Result<Box<dyn RuntimeIoWavePermit>, RuntimeIoWaveError>;
199}
200
201#[derive(Debug, Clone, Copy, PartialEq, Eq)]
202pub enum RuntimeIoWaveError {
203    Stopped(RuntimeCancellationReason),
204    ReservationExceeded {
205        requested_slots: usize,
206        reserved_slots: usize,
207    },
208}
209
210#[derive(Debug)]
211pub enum RuntimeIoWaveTryAcquire {
212    /// Capacity is available. Ungoverned contexts carry no permit.
213    Acquired(Option<Box<dyn RuntimeIoWavePermit>>),
214    /// Capacity is currently saturated and the caller should yield.
215    Pending,
216}
217
218impl Display for RuntimeIoWaveError {
219    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
220        match self {
221            Self::Stopped(reason) => write!(formatter, "runtime I/O wave stopped: {reason}"),
222            Self::ReservationExceeded {
223                requested_slots,
224                reserved_slots,
225            } => write!(
226                formatter,
227                "runtime I/O wave requested {requested_slots} slots, exceeding the {reserved_slots}-slot reservation"
228            ),
229        }
230    }
231}
232
233impl Error for RuntimeIoWaveError {
234    fn source(&self) -> Option<&(dyn Error + 'static)> {
235        match self {
236            Self::Stopped(reason) => Some(reason),
237            Self::ReservationExceeded { .. } => None,
238        }
239    }
240}
241
242#[derive(Debug, Clone)]
243pub struct RuntimeTaskContext {
244    cancellation: RuntimeCancellationToken,
245    deadline: Option<Instant>,
246    admitted_parallelism: NonZeroUsize,
247    executor_thread_limit: Option<NonZeroUsize>,
248    memory_reservation: Option<RuntimeMemoryReservation>,
249    io_wave_controller: Option<Arc<dyn RuntimeIoWaveController>>,
250}
251
252/// Memory already reserved for one task by the runtime governor.
253#[derive(Debug, Clone, Copy, PartialEq, Eq)]
254pub struct RuntimeMemoryReservation {
255    memory_bytes: u64,
256    result_bytes: u64,
257}
258
259impl RuntimeMemoryReservation {
260    pub const fn new(memory_bytes: u64, result_bytes: u64) -> Self {
261        Self {
262            memory_bytes,
263            result_bytes,
264        }
265    }
266
267    pub const fn memory_bytes(self) -> u64 {
268        self.memory_bytes
269    }
270
271    pub const fn result_bytes(self) -> u64 {
272        self.result_bytes
273    }
274
275    pub const fn intersect(self, other: Self) -> Self {
276        Self {
277            memory_bytes: if self.memory_bytes < other.memory_bytes {
278                self.memory_bytes
279            } else {
280                other.memory_bytes
281            },
282            result_bytes: if self.result_bytes < other.result_bytes {
283                self.result_bytes
284            } else {
285                other.result_bytes
286            },
287        }
288    }
289}
290
291impl RuntimeTaskContext {
292    pub fn new(cancellation: RuntimeCancellationToken, deadline: Option<Instant>) -> Self {
293        Self {
294            cancellation,
295            deadline,
296            admitted_parallelism: NonZeroUsize::MIN,
297            executor_thread_limit: None,
298            memory_reservation: None,
299            io_wave_controller: None,
300        }
301    }
302
303    pub fn without_deadline(cancellation: RuntimeCancellationToken) -> Self {
304        Self::new(cancellation, None)
305    }
306
307    pub fn with_timeout(timeout: Duration) -> Self {
308        Self::new(
309            RuntimeCancellationToken::new(),
310            Instant::now().checked_add(timeout),
311        )
312    }
313
314    pub fn cancellation(&self) -> &RuntimeCancellationToken {
315        &self.cancellation
316    }
317
318    pub fn child(&self) -> Self {
319        Self {
320            cancellation: self.cancellation.child(),
321            deadline: self.deadline,
322            admitted_parallelism: self.admitted_parallelism,
323            executor_thread_limit: self.executor_thread_limit,
324            memory_reservation: self.memory_reservation,
325            io_wave_controller: self.io_wave_controller.clone(),
326        }
327    }
328
329    /// Carries the CPU parallelism already reserved by the runtime governor.
330    ///
331    /// This is an execution ceiling, not a request to create worker threads.
332    /// Executors must still apply their operator memory and input-size limits.
333    pub fn with_admitted_parallelism(mut self, admitted_parallelism: NonZeroUsize) -> Self {
334        self.admitted_parallelism = admitted_parallelism;
335        self
336    }
337
338    pub fn admitted_parallelism(&self) -> NonZeroUsize {
339        self.admitted_parallelism
340    }
341
342    /// Carries the physical executor-thread ceiling derived by the runtime governor.
343    ///
344    /// This is distinct from [`Self::admitted_parallelism`], which is the CPU
345    /// reservation for one task. An absent limit denotes an ungoverned library call.
346    pub fn with_executor_thread_limit(mut self, limit: NonZeroUsize) -> Self {
347        self.executor_thread_limit = Some(
348            self.executor_thread_limit
349                .map_or(limit, |current| current.min(limit)),
350        );
351        self
352    }
353
354    pub fn executor_thread_limit(&self) -> Option<NonZeroUsize> {
355        self.executor_thread_limit
356    }
357
358    /// Carries the memory already reserved by the runtime governor.
359    ///
360    /// Governed executors must use this reservation instead of a static
361    /// per-query configuration limit. An absent reservation denotes an
362    /// ungoverned library call and preserves the configured fallback.
363    pub fn with_memory_reservation(mut self, reservation: RuntimeMemoryReservation) -> Self {
364        self.memory_reservation = Some(reservation);
365        self
366    }
367
368    pub fn memory_reservation(&self) -> Option<RuntimeMemoryReservation> {
369        self.memory_reservation
370    }
371
372    /// Binds the controller owned by a successful runtime admission.
373    pub fn with_io_wave_controller(mut self, controller: Arc<dyn RuntimeIoWaveController>) -> Self {
374        self.io_wave_controller = Some(controller);
375        self
376    }
377
378    /// Acquires capacity for one storage I/O wave when this is an admitted context.
379    ///
380    /// Raw contexts have no controller and return `Ok(None)`. The returned
381    /// permit must remain live only while physical reads are in flight.
382    pub fn acquire_io_wave(
383        &self,
384        slots: NonZeroUsize,
385    ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
386        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
387        self.io_wave_controller
388            .as_ref()
389            .map(|controller| controller.acquire(slots, self))
390            .transpose()
391    }
392
393    /// Attempts to acquire capacity for one storage I/O wave without blocking.
394    ///
395    /// Raw, ungoverned contexts return `Acquired(None)`. Governed contexts
396    /// return `Pending` while another wave owns the requested capacity.
397    pub fn try_acquire_io_wave(
398        &self,
399        slots: NonZeroUsize,
400    ) -> Result<RuntimeIoWaveTryAcquire, RuntimeIoWaveError> {
401        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
402        let Some(controller) = &self.io_wave_controller else {
403            return Ok(RuntimeIoWaveTryAcquire::Acquired(None));
404        };
405        controller.try_acquire(slots, self).map(|permit| {
406            permit.map_or(RuntimeIoWaveTryAcquire::Pending, |permit| {
407                RuntimeIoWaveTryAcquire::Acquired(Some(permit))
408            })
409        })
410    }
411
412    pub fn deadline(&self) -> Option<Instant> {
413        self.deadline
414    }
415
416    pub fn remaining(&self) -> Option<Duration> {
417        self.deadline
418            .map(|deadline| deadline.saturating_duration_since(Instant::now()))
419    }
420
421    pub fn checkpoint(&self) -> Result<(), RuntimeCancellationReason> {
422        if self.cancellation.is_cancelled() {
423            return Err(RuntimeCancellationReason::Cancelled);
424        }
425        if self
426            .deadline
427            .is_some_and(|deadline| Instant::now() >= deadline)
428        {
429            return Err(RuntimeCancellationReason::DeadlineExceeded);
430        }
431        Ok(())
432    }
433}
434
435impl Default for RuntimeTaskContext {
436    fn default() -> Self {
437        Self::without_deadline(RuntimeCancellationToken::new())
438    }
439}
440
441#[derive(Debug, Clone, Copy, PartialEq, Eq)]
442pub enum RuntimeCancellationReason {
443    Cancelled,
444    DeadlineExceeded,
445}
446
447impl RuntimeCancellationReason {
448    pub const fn as_str(self) -> &'static str {
449        match self {
450            Self::Cancelled => "cancelled",
451            Self::DeadlineExceeded => "deadline_exceeded",
452        }
453    }
454}
455
456impl Display for RuntimeCancellationReason {
457    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
458        formatter.write_str(self.as_str())
459    }
460}
461
462impl Error for RuntimeCancellationReason {}
463
464#[cfg(test)]
465mod tests {
466    use super::*;
467    use std::sync::atomic::AtomicUsize;
468    use std::sync::atomic::Ordering;
469    use std::task::Wake;
470
471    #[derive(Debug)]
472    struct RecordingIoController {
473        active_slots: Arc<AtomicUsize>,
474    }
475
476    #[derive(Debug)]
477    struct RecordingIoPermit {
478        active_slots: Arc<AtomicUsize>,
479        slots: usize,
480    }
481
482    impl Drop for RecordingIoPermit {
483        fn drop(&mut self) {
484            self.active_slots.fetch_sub(self.slots, Ordering::AcqRel);
485        }
486    }
487
488    impl RuntimeIoWaveController for RecordingIoController {
489        fn try_acquire(
490            &self,
491            slots: NonZeroUsize,
492            context: &RuntimeTaskContext,
493        ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
494            self.acquire(slots, context).map(Some)
495        }
496
497        fn acquire(
498            &self,
499            slots: NonZeroUsize,
500            context: &RuntimeTaskContext,
501        ) -> Result<Box<dyn RuntimeIoWavePermit>, RuntimeIoWaveError> {
502            context.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
503            self.active_slots.fetch_add(slots.get(), Ordering::AcqRel);
504            Ok(Box::new(RecordingIoPermit {
505                active_slots: Arc::clone(&self.active_slots),
506                slots: slots.get(),
507            }))
508        }
509    }
510
511    #[test]
512    fn cancellation_is_shared_across_context_clones() {
513        let token = RuntimeCancellationToken::new();
514        let context = RuntimeTaskContext::without_deadline(token.clone());
515        assert!(context.checkpoint().is_ok());
516        assert!(token.cancel());
517        assert!(!token.cancel());
518        assert_eq!(
519            context.checkpoint(),
520            Err(RuntimeCancellationReason::Cancelled)
521        );
522    }
523
524    #[test]
525    fn expired_deadline_fails_at_a_cooperative_checkpoint() {
526        let context = RuntimeTaskContext::new(
527            RuntimeCancellationToken::new(),
528            Some(Instant::now() - Duration::from_millis(1)),
529        );
530        assert_eq!(
531            context.checkpoint(),
532            Err(RuntimeCancellationReason::DeadlineExceeded)
533        );
534    }
535
536    #[test]
537    fn child_cancellation_is_local_and_parent_cancellation_propagates() {
538        let parent_token = RuntimeCancellationToken::new();
539        let parent = RuntimeTaskContext::without_deadline(parent_token.clone());
540        let child = parent.child();
541
542        assert!(child.cancellation().cancel());
543        assert!(child.checkpoint().is_err());
544        assert!(parent.checkpoint().is_ok());
545
546        let sibling = parent.child();
547        assert!(parent_token.cancel());
548        assert!(parent.checkpoint().is_err());
549        assert!(sibling.checkpoint().is_err());
550    }
551
552    #[derive(Default)]
553    struct WakeCounter(AtomicUsize);
554
555    impl Wake for WakeCounter {
556        fn wake(self: Arc<Self>) {
557            self.0.fetch_add(1, Ordering::AcqRel);
558        }
559    }
560
561    #[test]
562    fn cancellation_future_is_woken_by_parent_cancellation() {
563        let parent = RuntimeCancellationToken::new();
564        let child = parent.child();
565        let counter = Arc::new(WakeCounter::default());
566        let waker = Waker::from(Arc::clone(&counter));
567        let mut context = Context::from_waker(&waker);
568        let mut cancelled = Box::pin(child.cancelled());
569
570        assert!(cancelled.as_mut().poll(&mut context).is_pending());
571        assert!(parent.cancel());
572        assert_eq!(counter.0.load(Ordering::Acquire), 1);
573        assert!(cancelled.as_mut().poll(&mut context).is_ready());
574    }
575
576    #[test]
577    fn dropped_cancellation_future_unregisters_from_its_parent_chain() {
578        let parent = RuntimeCancellationToken::new();
579        let child = parent.child();
580        let counter = Arc::new(WakeCounter::default());
581        let waker = Waker::from(Arc::clone(&counter));
582        let mut context = Context::from_waker(&waker);
583        let mut cancelled = Box::pin(child.cancelled());
584
585        assert!(cancelled.as_mut().poll(&mut context).is_pending());
586        drop(cancelled);
587        assert!(parent.cancel());
588        assert_eq!(counter.0.load(Ordering::Acquire), 0);
589    }
590
591    #[test]
592    fn child_preserves_admitted_parallelism() {
593        let context = RuntimeTaskContext::default()
594            .with_admitted_parallelism(NonZeroUsize::new(4).expect("test parallelism is non-zero"))
595            .with_executor_thread_limit(
596                NonZeroUsize::new(2).expect("test thread limit is non-zero"),
597            );
598
599        assert_eq!(context.admitted_parallelism().get(), 4);
600        assert_eq!(context.child().admitted_parallelism().get(), 4);
601        assert_eq!(context.child().executor_thread_limit().unwrap().get(), 2);
602    }
603
604    #[test]
605    fn child_preserves_memory_reservation() {
606        let reservation = RuntimeMemoryReservation::new(128, 32);
607        let context = RuntimeTaskContext::default().with_memory_reservation(reservation);
608
609        assert_eq!(context.memory_reservation(), Some(reservation));
610        assert_eq!(context.child().memory_reservation(), Some(reservation));
611    }
612
613    #[test]
614    fn child_preserves_io_controller_and_permit_releases_slots() {
615        let active_slots = Arc::new(AtomicUsize::new(0));
616        let context = RuntimeTaskContext::default().with_io_wave_controller(Arc::new(
617            RecordingIoController {
618                active_slots: Arc::clone(&active_slots),
619            },
620        ));
621
622        let permit = context
623            .child()
624            .acquire_io_wave(NonZeroUsize::new(2).unwrap())
625            .unwrap()
626            .unwrap();
627        assert_eq!(active_slots.load(Ordering::Acquire), 2);
628        drop(permit);
629        assert_eq!(active_slots.load(Ordering::Acquire), 0);
630    }
631
632    #[test]
633    fn ungoverned_try_acquire_is_immediately_ready() {
634        let acquired = RuntimeTaskContext::default()
635            .try_acquire_io_wave(NonZeroUsize::new(2).unwrap())
636            .unwrap();
637
638        assert!(matches!(acquired, RuntimeIoWaveTryAcquire::Acquired(None)));
639    }
640}
hawdb_core/cancellation.rs

hawdb_core/
cancellation.rs
