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 crate::time::{Duration, Instant};
16use std::collections::{BTreeMap, BTreeSet};
17use std::error::Error;
18use std::fmt::Debug;
19use std::fmt::{self, Display, Formatter};
20use std::future::Future;
21use std::num::NonZeroUsize;
22use std::pin::Pin;
23use std::sync::atomic::{AtomicU64, AtomicU8, Ordering};
24use std::sync::{Arc, Mutex, Weak};
25use std::task::{Context, Poll, Waker};
26
27const NOT_CANCELLED: u8 = 0;
28const CANCELLED: u8 = 1;
29const RESOURCE_EXHAUSTED: u8 = 2;
30
31#[derive(Debug, Clone)]
32pub struct RuntimeCancellationToken {
33    state: Arc<RuntimeCancellationState>,
34}
35
36#[derive(Debug)]
37struct RuntimeCancellationState {
38    cancellation_reason: AtomicU8,
39    parents: [Option<RuntimeCancellationToken>; 2],
40    next_waiter_id: AtomicU64,
41    waiters: Mutex<BTreeMap<u64, Waker>>,
42}
43
44impl Default for RuntimeCancellationToken {
45    fn default() -> Self {
46        Self {
47            state: Arc::new(RuntimeCancellationState {
48                cancellation_reason: AtomicU8::new(NOT_CANCELLED),
49                parents: [None, None],
50                next_waiter_id: AtomicU64::new(1),
51                waiters: Mutex::new(BTreeMap::new()),
52            }),
53        }
54    }
55}
56
57impl RuntimeCancellationToken {
58    pub fn new() -> Self {
59        Self::default()
60    }
61
62    pub fn cancel(&self) -> bool {
63        self.cancel_with_code(CANCELLED)
64    }
65
66    /// Stops cooperating siblings after an unrecoverable resource refusal.
67    /// The first local cancellation cause is retained. Cancelling a child does
68    /// not cancel its parent; the operation must still return its original error.
69    pub fn cancel_for_resource_exhaustion(&self) -> bool {
70        self.cancel_with_code(RESOURCE_EXHAUSTED)
71    }
72
73    fn cancel_with_code(&self, reason: u8) -> bool {
74        if self
75            .state
76            .cancellation_reason
77            .compare_exchange(NOT_CANCELLED, reason, Ordering::AcqRel, Ordering::Acquire)
78            .is_err()
79        {
80            return false;
81        }
82        let waiters = {
83            let mut waiters = self
84                .state
85                .waiters
86                .lock()
87                .unwrap_or_else(|poisoned| poisoned.into_inner());
88            std::mem::take(&mut *waiters)
89        };
90        for (_, waiter) in waiters {
91            waiter.wake();
92        }
93        true
94    }
95
96    pub fn is_cancelled(&self) -> bool {
97        self.reason().is_some()
98    }
99
100    fn reason(&self) -> Option<RuntimeCancellationReason> {
101        match self.state.cancellation_reason.load(Ordering::Acquire) {
102            CANCELLED => Some(RuntimeCancellationReason::Cancelled),
103            RESOURCE_EXHAUSTED => Some(RuntimeCancellationReason::ResourceExhausted),
104            _ => self
105                .state
106                .parents
107                .iter()
108                .flatten()
109                .find_map(RuntimeCancellationToken::reason),
110        }
111    }
112
113    pub fn child(&self) -> Self {
114        Self {
115            state: Arc::new(RuntimeCancellationState {
116                cancellation_reason: AtomicU8::new(NOT_CANCELLED),
117                parents: [Some(self.clone()), None],
118                next_waiter_id: AtomicU64::new(1),
119                waiters: Mutex::new(BTreeMap::new()),
120            }),
121        }
122    }
123
124    fn linked(&self, other: &Self) -> Self {
125        if Arc::ptr_eq(&self.state, &other.state) {
126            return self.child();
127        }
128        Self {
129            state: Arc::new(RuntimeCancellationState {
130                cancellation_reason: AtomicU8::new(NOT_CANCELLED),
131                parents: [Some(self.clone()), Some(other.clone())],
132                next_waiter_id: AtomicU64::new(1),
133                waiters: Mutex::new(BTreeMap::new()),
134            }),
135        }
136    }
137
138    /// Resolves when this token or any parent token is cancelled.
139    pub fn cancelled(&self) -> RuntimeCancellationFuture<'_> {
140        RuntimeCancellationFuture {
141            token: self,
142            registrations: Vec::new(),
143        }
144    }
145}
146
147struct RuntimeCancellationRegistration {
148    state: Weak<RuntimeCancellationState>,
149    waiter_id: u64,
150}
151
152/// A runtime-neutral future returned by [`RuntimeCancellationToken::cancelled`].
153pub struct RuntimeCancellationFuture<'a> {
154    token: &'a RuntimeCancellationToken,
155    registrations: Vec<RuntimeCancellationRegistration>,
156}
157
158impl Future for RuntimeCancellationFuture<'_> {
159    type Output = ();
160
161    fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
162        if self.token.is_cancelled() {
163            return Poll::Ready(());
164        }
165
166        if self.registrations.is_empty() {
167            let mut pending = vec![self.token];
168            let mut visited = BTreeSet::new();
169            while let Some(current) = pending.pop() {
170                if !visited.insert(Arc::as_ptr(&current.state)) {
171                    continue;
172                }
173                let waiter_id = current.state.next_waiter_id.fetch_add(1, Ordering::Relaxed);
174                current
175                    .state
176                    .waiters
177                    .lock()
178                    .unwrap_or_else(|poisoned| poisoned.into_inner())
179                    .insert(waiter_id, context.waker().clone());
180                self.registrations.push(RuntimeCancellationRegistration {
181                    state: Arc::downgrade(&current.state),
182                    waiter_id,
183                });
184                pending.extend(current.state.parents.iter().flatten());
185            }
186        } else {
187            for registration in &self.registrations {
188                if let Some(state) = registration.state.upgrade()
189                    && let Some(waiter) = state
190                        .waiters
191                        .lock()
192                        .unwrap_or_else(|poisoned| poisoned.into_inner())
193                        .get_mut(&registration.waiter_id)
194                    && !waiter.will_wake(context.waker())
195                {
196                    *waiter = context.waker().clone();
197                }
198            }
199        }
200
201        if self.token.is_cancelled() {
202            Poll::Ready(())
203        } else {
204            Poll::Pending
205        }
206    }
207}
208
209impl Drop for RuntimeCancellationFuture<'_> {
210    fn drop(&mut self) {
211        for registration in &self.registrations {
212            if let Some(state) = registration.state.upgrade() {
213                state
214                    .waiters
215                    .lock()
216                    .unwrap_or_else(|poisoned| poisoned.into_inner())
217                    .remove(&registration.waiter_id);
218            }
219        }
220    }
221}
222
223/// A drop guard for one admitted storage I/O wave.
224pub trait RuntimeIoWavePermit: Debug + Send {}
225
226impl<T: Debug + Send> RuntimeIoWavePermit for T {}
227
228/// Execution-owned hook for acquiring the I/O capacity declared at admission.
229pub trait RuntimeIoWaveController: Debug + Send + Sync {
230    /// Attempts to acquire one I/O wave without blocking the calling thread.
231    ///
232    /// `Ok(None)` means that the reservation is currently saturated. Async
233    /// runtime adapters must yield before retrying instead of spinning.
234    fn try_acquire(
235        &self,
236        slots: NonZeroUsize,
237        context: &RuntimeTaskContext,
238    ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError>;
239
240    fn acquire(
241        &self,
242        slots: NonZeroUsize,
243        context: &RuntimeTaskContext,
244    ) -> Result<Box<dyn RuntimeIoWavePermit>, RuntimeIoWaveError>;
245}
246
247#[derive(Debug, Clone, Copy, PartialEq, Eq)]
248pub enum RuntimeIoWaveError {
249    Stopped(RuntimeCancellationReason),
250    ReservationExceeded {
251        requested_slots: usize,
252        reserved_slots: usize,
253    },
254}
255
256#[derive(Debug)]
257pub enum RuntimeIoWaveTryAcquire {
258    /// Capacity is available. Ungoverned contexts carry no permit.
259    Acquired(Option<Box<dyn RuntimeIoWavePermit>>),
260    /// Capacity is currently saturated and the caller should yield.
261    Pending,
262}
263
264impl Display for RuntimeIoWaveError {
265    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
266        match self {
267            Self::Stopped(reason) => write!(formatter, "runtime I/O wave stopped: {reason}"),
268            Self::ReservationExceeded {
269                requested_slots,
270                reserved_slots,
271            } => write!(
272                formatter,
273                "runtime I/O wave requested {requested_slots} slots, exceeding the {reserved_slots}-slot reservation"
274            ),
275        }
276    }
277}
278
279impl Error for RuntimeIoWaveError {
280    fn source(&self) -> Option<&(dyn Error + 'static)> {
281        match self {
282            Self::Stopped(reason) => Some(reason),
283            Self::ReservationExceeded { .. } => None,
284        }
285    }
286}
287
288mod memory;
289mod work;
290pub use memory::{RuntimeMemoryController, RuntimeMemoryError, RuntimeMemoryPermit};
291
292#[derive(Debug, Clone)]
293pub struct RuntimeTaskContext {
294    cancellation: RuntimeCancellationToken,
295    deadline: Option<Instant>,
296    admitted_parallelism: NonZeroUsize,
297    executor_thread_limit: Option<NonZeroUsize>,
298    memory_reservation: Option<RuntimeMemoryReservation>,
299    memory_controller: Option<Arc<dyn RuntimeMemoryController>>,
300    io_wave_controller: Option<Arc<dyn RuntimeIoWaveController>>,
301    work_budget: Option<Arc<work::RuntimeWorkBudget>>,
302}
303
304/// Memory already reserved for one task by the runtime governor.
305#[derive(Debug, Clone, Copy, PartialEq, Eq)]
306pub struct RuntimeMemoryReservation {
307    memory_bytes: u64,
308    result_bytes: u64,
309}
310
311impl RuntimeMemoryReservation {
312    pub const fn new(memory_bytes: u64, result_bytes: u64) -> Self {
313        Self {
314            memory_bytes,
315            result_bytes,
316        }
317    }
318
319    pub const fn memory_bytes(self) -> u64 {
320        self.memory_bytes
321    }
322
323    pub const fn result_bytes(self) -> u64 {
324        self.result_bytes
325    }
326
327    pub const fn intersect(self, other: Self) -> Self {
328        Self {
329            memory_bytes: if self.memory_bytes < other.memory_bytes {
330                self.memory_bytes
331            } else {
332                other.memory_bytes
333            },
334            result_bytes: if self.result_bytes < other.result_bytes {
335                self.result_bytes
336            } else {
337                other.result_bytes
338            },
339        }
340    }
341}
342
343impl RuntimeTaskContext {
344    pub fn new(cancellation: RuntimeCancellationToken, deadline: Option<Instant>) -> Self {
345        Self {
346            cancellation,
347            deadline,
348            admitted_parallelism: NonZeroUsize::MIN,
349            executor_thread_limit: None,
350            memory_reservation: None,
351            memory_controller: None,
352            io_wave_controller: None,
353            work_budget: None,
354        }
355    }
356
357    pub fn without_deadline(cancellation: RuntimeCancellationToken) -> Self {
358        Self::new(cancellation, None)
359    }
360
361    pub fn with_timeout(timeout: Duration) -> Self {
362        Self::new(
363            RuntimeCancellationToken::new(),
364            Instant::now().checked_add(timeout),
365        )
366    }
367
368    pub fn cancellation(&self) -> &RuntimeCancellationToken {
369        &self.cancellation
370    }
371
372    pub fn child(&self) -> Self {
373        Self {
374            cancellation: self.cancellation.child(),
375            deadline: self.deadline,
376            admitted_parallelism: self.admitted_parallelism,
377            executor_thread_limit: self.executor_thread_limit,
378            memory_reservation: self.memory_reservation,
379            memory_controller: self.memory_controller.clone(),
380            io_wave_controller: self.io_wave_controller.clone(),
381            work_budget: self.work_budget.clone(),
382        }
383    }
384
385    /// Applies statement constraints without replacing an enclosing admission.
386    /// Cancellation from either scope stops the statement; cancelling the combined
387    /// scope does not cancel its parents. Budgets remain shared and can only tighten.
388    /// Independent controllers cannot represent one admission and are rejected.
389    pub fn constrained_by(&self, request: &Self) -> crate::Result<Self> {
390        fn controller<T: ?Sized>(
391            parent: &Option<Arc<T>>,
392            request: &Option<Arc<T>>,
393            resource: &str,
394        ) -> crate::Result<Option<Arc<T>>> {
395            match (parent, request) {
396                (Some(parent), Some(request)) if !Arc::ptr_eq(parent, request) => {
397                    Err(crate::HawDBError::Execution(format!(
398                        "request task context carries a different {resource} admission"
399                    )))
400                }
401                (Some(parent), _) => Ok(Some(parent.clone())),
402                (_, request) => Ok(request.clone()),
403            }
404        }
405        fn minimum<T: Ord + Copy>(left: Option<T>, right: Option<T>) -> Option<T> {
406            match (left, right) {
407                (Some(left), Some(right)) => Some(left.min(right)),
408                (left, right) => left.or(right),
409            }
410        }
411        Ok(Self {
412            cancellation: self.cancellation.linked(&request.cancellation),
413            deadline: minimum(self.deadline, request.deadline),
414            admitted_parallelism: self.admitted_parallelism.min(request.admitted_parallelism),
415            executor_thread_limit: minimum(
416                self.executor_thread_limit,
417                request.executor_thread_limit,
418            ),
419            memory_reservation: match (self.memory_reservation, request.memory_reservation) {
420                (Some(parent), Some(request)) => Some(parent.intersect(request)),
421                (parent, request) => parent.or(request),
422            },
423            memory_controller: controller(
424                &self.memory_controller,
425                &request.memory_controller,
426                "memory",
427            )?,
428            io_wave_controller: controller(
429                &self.io_wave_controller,
430                &request.io_wave_controller,
431                "I/O",
432            )?,
433            work_budget: work::RuntimeWorkBudget::intersect(
434                &self.work_budget,
435                &request.work_budget,
436            ),
437        })
438    }
439
440    /// Carries the CPU parallelism already reserved by the runtime governor.
441    ///
442    /// This is an execution ceiling, not a request to create worker threads.
443    /// Executors must still apply their operator memory and input-size limits.
444    pub fn with_admitted_parallelism(mut self, admitted_parallelism: NonZeroUsize) -> Self {
445        self.admitted_parallelism = admitted_parallelism;
446        self
447    }
448
449    pub fn admitted_parallelism(&self) -> NonZeroUsize {
450        self.admitted_parallelism
451    }
452
453    /// Carries the physical executor-thread ceiling derived by the runtime governor.
454    ///
455    /// This is distinct from [`Self::admitted_parallelism`], which is the CPU
456    /// reservation for one task. An absent limit denotes an ungoverned library call.
457    pub fn with_executor_thread_limit(mut self, limit: NonZeroUsize) -> Self {
458        self.executor_thread_limit = Some(
459            self.executor_thread_limit
460                .map_or(limit, |current| current.min(limit)),
461        );
462        self
463    }
464
465    pub fn executor_thread_limit(&self) -> Option<NonZeroUsize> {
466        self.executor_thread_limit
467    }
468
469    /// Carries the memory already reserved by the runtime governor.
470    ///
471    /// Governed executors must use this reservation instead of a static
472    /// per-query configuration limit. An absent reservation denotes an
473    /// ungoverned library call and preserves the configured fallback.
474    pub fn with_memory_reservation(mut self, reservation: RuntimeMemoryReservation) -> Self {
475        self.memory_reservation = Some(reservation);
476        self
477    }
478
479    pub fn memory_reservation(&self) -> Option<RuntimeMemoryReservation> {
480        self.memory_reservation
481    }
482
483    /// Whether execution carries dynamic resource controllers rather than only
484    /// a static ceiling. Synchronous adapters must preserve these controllers
485    /// and must not replace an originating governor admission with a local cap.
486    pub fn has_resource_controllers(&self) -> bool {
487        self.memory_controller.is_some() || self.io_wave_controller.is_some()
488    }
489
490    /// Binds actual allocation ownership to an already admitted task.
491    pub fn with_memory_controller(mut self, controller: Arc<dyn RuntimeMemoryController>) -> Self {
492        self.memory_controller = Some(controller);
493        self
494    }
495
496    /// Reserve before allocating and retain the returned lease until destruction.
497    /// Clones and children share their controller. Ungoverned contexts return
498    /// `None`; a supplied static ceiling still limits an individual allocation.
499    pub fn reserve_working_memory(
500        &self,
501        bytes: u64,
502    ) -> Result<Option<Box<dyn RuntimeMemoryPermit>>, RuntimeMemoryError> {
503        self.checkpoint().map_err(RuntimeMemoryError::Stopped)?;
504        if let Some(reservation) = self.memory_reservation
505            && bytes > reservation.memory_bytes()
506        {
507            return Err(RuntimeMemoryError::ReservationExceeded {
508                requested_bytes: bytes,
509                available_bytes: reservation.memory_bytes(),
510            });
511        }
512        self.memory_controller
513            .as_ref()
514            .map(|controller| {
515                controller.reserve(
516                    bytes,
517                    self.memory_reservation
518                        .map_or(u64::MAX, RuntimeMemoryReservation::memory_bytes),
519                )
520            })
521            .transpose()
522    }
523
524    /// Binds the controller owned by a successful runtime admission.
525    pub fn with_io_wave_controller(mut self, controller: Arc<dyn RuntimeIoWaveController>) -> Self {
526        self.io_wave_controller = Some(controller);
527        self
528    }
529
530    /// Acquires capacity for one storage I/O wave when this is an admitted context.
531    ///
532    /// Raw contexts have no controller and return `Ok(None)`. The returned
533    /// permit must remain live only while physical reads are in flight.
534    pub fn acquire_io_wave(
535        &self,
536        slots: NonZeroUsize,
537    ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
538        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
539        self.io_wave_controller
540            .as_ref()
541            .map(|controller| controller.acquire(slots, self))
542            .transpose()
543    }
544
545    /// Attempts to acquire capacity for one storage I/O wave without blocking.
546    ///
547    /// Raw, ungoverned contexts return `Acquired(None)`. Governed contexts
548    /// return `Pending` while another wave owns the requested capacity.
549    pub fn try_acquire_io_wave(
550        &self,
551        slots: NonZeroUsize,
552    ) -> Result<RuntimeIoWaveTryAcquire, RuntimeIoWaveError> {
553        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
554        let Some(controller) = &self.io_wave_controller else {
555            return Ok(RuntimeIoWaveTryAcquire::Acquired(None));
556        };
557        controller.try_acquire(slots, self).map(|permit| {
558            permit.map_or(RuntimeIoWaveTryAcquire::Pending, |permit| {
559                RuntimeIoWaveTryAcquire::Acquired(Some(permit))
560            })
561        })
562    }
563
564    /// Starts a bounded execution scope while retaining every parent limit and
565    /// the original cancellation, memory and I/O admission. Clones share usage.
566    pub fn with_work_limit(mut self, limit: usize) -> Self {
567        self.work_budget = Some(Arc::new(work::RuntimeWorkBudget::new(
568            limit,
569            self.work_budget.take(),
570        )));
571        self
572    }
573
574    pub fn has_work_limit(&self) -> bool {
575        self.work_budget.is_some()
576    }
577
578    pub fn work_units(&self) -> usize {
579        self.work_budget.as_ref().map_or(0, |budget| budget.used())
580    }
581
582    /// Preserve a typed refusal across adapters whose error transport is text.
583    pub fn check_work_limit(&self) -> crate::Result<()> {
584        match self
585            .work_budget
586            .as_ref()
587            .and_then(|budget| budget.refusal())
588        {
589            Some(error) => Err(error.into()),
590            None => Ok(()),
591        }
592    }
593
594    /// Charge explicit candidate work before examining it, including rejected
595    /// candidates. Checkpoints and allocation byte counts are not work units.
596    pub fn charge_work(&self, units: usize) -> crate::Result<()> {
597        let Some(budget) = &self.work_budget else {
598            return Ok(());
599        };
600        self.checkpoint().map_err(|reason| {
601            crate::HawDBError::Execution(format!("runtime task stopped: {reason}"))
602        })?;
603        budget.charge(units)
604    }
605
606    pub fn deadline(&self) -> Option<Instant> {
607        self.deadline
608    }
609
610    pub fn remaining(&self) -> Option<Duration> {
611        self.deadline
612            .map(|deadline| deadline.saturating_duration_since(Instant::now()))
613    }
614
615    pub fn checkpoint(&self) -> Result<(), RuntimeCancellationReason> {
616        if let Some(reason) = self.cancellation.reason() {
617            return Err(reason);
618        }
619        if self
620            .deadline
621            .is_some_and(|deadline| Instant::now() >= deadline)
622        {
623            return Err(RuntimeCancellationReason::DeadlineExceeded);
624        }
625        Ok(())
626    }
627}
628
629impl Default for RuntimeTaskContext {
630    fn default() -> Self {
631        Self::without_deadline(RuntimeCancellationToken::new())
632    }
633}
634
635#[derive(Debug, Clone, Copy, PartialEq, Eq)]
636pub enum RuntimeCancellationReason {
637    Cancelled,
638    DeadlineExceeded,
639    ResourceExhausted,
640}
641
642impl RuntimeCancellationReason {
643    pub const fn as_str(self) -> &'static str {
644        match self {
645            Self::Cancelled => "cancelled",
646            Self::DeadlineExceeded => "deadline_exceeded",
647            Self::ResourceExhausted => "resource_exhausted",
648        }
649    }
650}
651
652impl Display for RuntimeCancellationReason {
653    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
654        formatter.write_str(self.as_str())
655    }
656}
657
658impl Error for RuntimeCancellationReason {}
659
660#[cfg(test)]
661mod tests {
662    use super::*;
663    use std::sync::atomic::AtomicUsize;
664    use std::sync::atomic::Ordering;
665    use std::task::Wake;
666
667    #[derive(Debug)]
668    struct RecordingIoController {
669        active_slots: Arc<AtomicUsize>,
670    }
671
672    #[derive(Debug)]
673    struct RecordingIoPermit {
674        active_slots: Arc<AtomicUsize>,
675        slots: usize,
676    }
677
678    impl Drop for RecordingIoPermit {
679        fn drop(&mut self) {
680            self.active_slots.fetch_sub(self.slots, Ordering::AcqRel);
681        }
682    }
683
684    impl RuntimeIoWaveController for RecordingIoController {
685        fn try_acquire(
686            &self,
687            slots: NonZeroUsize,
688            context: &RuntimeTaskContext,
689        ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
690            self.acquire(slots, context).map(Some)
691        }
692
693        fn acquire(
694            &self,
695            slots: NonZeroUsize,
696            context: &RuntimeTaskContext,
697        ) -> Result<Box<dyn RuntimeIoWavePermit>, RuntimeIoWaveError> {
698            context.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
699            self.active_slots.fetch_add(slots.get(), Ordering::AcqRel);
700            Ok(Box::new(RecordingIoPermit {
701                active_slots: Arc::clone(&self.active_slots),
702                slots: slots.get(),
703            }))
704        }
705    }
706
707    #[test]
708    fn cancellation_is_shared_across_context_clones() {
709        let token = RuntimeCancellationToken::new();
710        let context = RuntimeTaskContext::without_deadline(token.clone());
711        assert!(context.checkpoint().is_ok());
712        assert!(token.cancel());
713        assert!(!token.cancel());
714        assert_eq!(
715            context.checkpoint(),
716            Err(RuntimeCancellationReason::Cancelled)
717        );
718    }
719
720    #[test]
721    fn expired_deadline_fails_at_a_cooperative_checkpoint() {
722        let context = RuntimeTaskContext::new(
723            RuntimeCancellationToken::new(),
724            Some(Instant::now() - Duration::from_millis(1)),
725        );
726        assert_eq!(
727            context.checkpoint(),
728            Err(RuntimeCancellationReason::DeadlineExceeded)
729        );
730    }
731
732    #[test]
733    fn child_cancellation_is_local_and_parent_cancellation_propagates() {
734        let parent_token = RuntimeCancellationToken::new();
735        let parent = RuntimeTaskContext::without_deadline(parent_token.clone());
736        let child = parent.child();
737
738        assert!(child.cancellation().cancel());
739        assert!(child.checkpoint().is_err());
740        assert!(parent.checkpoint().is_ok());
741
742        let sibling = parent.child();
743        assert!(parent_token.cancel());
744        assert!(parent.checkpoint().is_err());
745        assert!(sibling.checkpoint().is_err());
746    }
747
748    #[test]
749    fn resource_exhaustion_keeps_its_cause_and_is_local_to_the_child() {
750        let parent = RuntimeTaskContext::default();
751        let child = parent.child();
752        let clone = child.clone();
753        let grandchild = child.child();
754        assert!(child.cancellation().cancel_for_resource_exhaustion());
755        assert!(!clone.cancellation().cancel());
756        assert!(!clone.cancellation().cancel_for_resource_exhaustion());
757        for context in [&child, &clone, &grandchild] {
758            assert_eq!(
759                context.checkpoint(),
760                Err(RuntimeCancellationReason::ResourceExhausted)
761            );
762        }
763        assert_eq!(parent.checkpoint(), Ok(()));
764        assert_eq!(parent.child().checkpoint(), Ok(()));
765
766        let cancelled = parent.child();
767        assert!(cancelled.cancellation().cancel());
768        assert!(!cancelled.cancellation().cancel_for_resource_exhaustion());
769        assert_eq!(
770            cancelled.checkpoint(),
771            Err(RuntimeCancellationReason::Cancelled)
772        );
773    }
774
775    #[derive(Default)]
776    struct WakeCounter(AtomicUsize);
777
778    impl Wake for WakeCounter {
779        fn wake(self: Arc<Self>) {
780            self.0.fetch_add(1, Ordering::AcqRel);
781        }
782    }
783
784    #[test]
785    fn constrained_context_observes_both_cancellations_without_cancelling_parents() {
786        for cancel_parent in [false, true] {
787            let parent = RuntimeTaskContext::default();
788            let request = RuntimeTaskContext::default();
789            let combined = parent.constrained_by(&request).unwrap();
790            let counter = Arc::new(WakeCounter::default());
791            let waker = Waker::from(counter.clone());
792            let mut context = Context::from_waker(&waker);
793            let mut cancelled = Box::pin(combined.cancellation().cancelled());
794            assert!(cancelled.as_mut().poll(&mut context).is_pending());
795            if cancel_parent {
796                parent.cancellation().cancel();
797            } else {
798                request.cancellation().cancel();
799            }
800            assert_eq!(counter.0.load(Ordering::Acquire), 1);
801            assert!(cancelled.as_mut().poll(&mut context).is_ready());
802            assert_eq!(
803                combined.checkpoint(),
804                Err(RuntimeCancellationReason::Cancelled)
805            );
806        }
807        let parent = RuntimeTaskContext::default();
808        let request = RuntimeTaskContext::default();
809        let combined = parent.constrained_by(&request).unwrap();
810        combined.cancellation().cancel();
811        assert!(parent.checkpoint().is_ok());
812        assert!(request.checkpoint().is_ok());
813    }
814
815    #[test]
816    fn constrained_descendants_register_shared_ancestors_once_and_unregister_on_drop() {
817        let parent = RuntimeTaskContext::default();
818        let request = parent.child();
819        let combined = parent.constrained_by(&request).unwrap();
820        let waker = Waker::from(Arc::new(WakeCounter::default()));
821        let mut context = Context::from_waker(&waker);
822        let mut cancelled = Box::pin(combined.cancellation().cancelled());
823        assert!(cancelled.as_mut().poll(&mut context).is_pending());
824        assert_eq!(parent.cancellation.state.waiters.lock().unwrap().len(), 1);
825        assert_eq!(request.cancellation.state.waiters.lock().unwrap().len(), 1);
826        drop(cancelled);
827        assert!(parent.cancellation.state.waiters.lock().unwrap().is_empty());
828        assert!(request
829            .cancellation
830            .state
831            .waiters
832            .lock()
833            .unwrap()
834            .is_empty());
835    }
836
837    #[test]
838    fn constrained_context_tightens_limits_and_keeps_the_original_io_controller() {
839        let active_slots = Arc::new(AtomicUsize::new(0));
840        let parent = RuntimeTaskContext::with_timeout(Duration::from_secs(5))
841            .with_memory_reservation(RuntimeMemoryReservation::new(100, 20))
842            .with_admitted_parallelism(NonZeroUsize::new(4).unwrap())
843            .with_io_wave_controller(Arc::new(RecordingIoController {
844                active_slots: active_slots.clone(),
845            }));
846        let request = RuntimeTaskContext::with_timeout(Duration::from_secs(10))
847            .with_memory_reservation(RuntimeMemoryReservation::new(30, 50))
848            .with_admitted_parallelism(NonZeroUsize::new(2).unwrap());
849        let combined = parent.constrained_by(&request).unwrap();
850        assert_eq!(combined.deadline(), parent.deadline());
851        assert_eq!(
852            combined.memory_reservation(),
853            Some(RuntimeMemoryReservation::new(30, 20))
854        );
855        assert_eq!(combined.admitted_parallelism().get(), 2);
856        let io = combined.acquire_io_wave(NonZeroUsize::MIN).unwrap();
857        assert_eq!(active_slots.load(Ordering::Acquire), 1);
858        drop(io);
859        assert_eq!(active_slots.load(Ordering::Acquire), 0);
860        assert!(parent.constrained_by(&parent.child()).is_ok());
861        let independent =
862            request.with_io_wave_controller(Arc::new(RecordingIoController { active_slots }));
863        assert!(parent.constrained_by(&independent).is_err());
864    }
865
866    #[test]
867    fn cancellation_future_is_woken_by_parent_cancellation() {
868        let parent = RuntimeCancellationToken::new();
869        let child = parent.child();
870        let counter = Arc::new(WakeCounter::default());
871        let waker = Waker::from(Arc::clone(&counter));
872        let mut context = Context::from_waker(&waker);
873        let mut cancelled = Box::pin(child.cancelled());
874
875        assert!(cancelled.as_mut().poll(&mut context).is_pending());
876        assert!(parent.cancel());
877        assert_eq!(counter.0.load(Ordering::Acquire), 1);
878        assert!(cancelled.as_mut().poll(&mut context).is_ready());
879    }
880
881    #[test]
882    fn resource_exhaustion_wakes_waiters_on_descendants() {
883        let parent = RuntimeTaskContext::default();
884        let child = parent.child();
885        let counter = Arc::new(WakeCounter::default());
886        let waker = Waker::from(Arc::clone(&counter));
887        let mut context = Context::from_waker(&waker);
888        let mut cancelled = Box::pin(child.cancellation().cancelled());
889        assert!(cancelled.as_mut().poll(&mut context).is_pending());
890        assert!(parent.cancellation().cancel_for_resource_exhaustion());
891        assert_eq!(counter.0.load(Ordering::Acquire), 1);
892        assert!(cancelled.as_mut().poll(&mut context).is_ready());
893        assert_eq!(
894            child.checkpoint(),
895            Err(RuntimeCancellationReason::ResourceExhausted)
896        );
897    }
898
899    #[test]
900    fn dropped_cancellation_future_unregisters_from_its_parent_chain() {
901        let parent = RuntimeCancellationToken::new();
902        let child = parent.child();
903        let counter = Arc::new(WakeCounter::default());
904        let waker = Waker::from(Arc::clone(&counter));
905        let mut context = Context::from_waker(&waker);
906        let mut cancelled = Box::pin(child.cancelled());
907
908        assert!(cancelled.as_mut().poll(&mut context).is_pending());
909        drop(cancelled);
910        assert!(parent.cancel());
911        assert_eq!(counter.0.load(Ordering::Acquire), 0);
912    }
913
914    #[test]
915    fn child_preserves_admitted_parallelism() {
916        let context = RuntimeTaskContext::default()
917            .with_admitted_parallelism(NonZeroUsize::new(4).expect("test parallelism is non-zero"))
918            .with_executor_thread_limit(
919                NonZeroUsize::new(2).expect("test thread limit is non-zero"),
920            );
921
922        assert_eq!(context.admitted_parallelism().get(), 4);
923        assert_eq!(context.child().admitted_parallelism().get(), 4);
924        assert_eq!(context.child().executor_thread_limit().unwrap().get(), 2);
925    }
926
927    #[test]
928    fn child_preserves_memory_reservation() {
929        let reservation = RuntimeMemoryReservation::new(128, 32);
930        let context = RuntimeTaskContext::default().with_memory_reservation(reservation);
931
932        assert_eq!(context.memory_reservation(), Some(reservation));
933        assert_eq!(context.child().memory_reservation(), Some(reservation));
934    }
935
936    #[test]
937    fn child_preserves_io_controller_and_permit_releases_slots() {
938        let active_slots = Arc::new(AtomicUsize::new(0));
939        let context = RuntimeTaskContext::default().with_io_wave_controller(Arc::new(
940            RecordingIoController {
941                active_slots: Arc::clone(&active_slots),
942            },
943        ));
944
945        let permit = context
946            .child()
947            .acquire_io_wave(NonZeroUsize::new(2).unwrap())
948            .unwrap()
949            .unwrap();
950        assert_eq!(active_slots.load(Ordering::Acquire), 2);
951        drop(permit);
952        assert_eq!(active_slots.load(Ordering::Acquire), 0);
953    }
954
955    #[test]
956    fn ungoverned_try_acquire_is_immediately_ready() {
957        let acquired = RuntimeTaskContext::default()
958            .try_acquire_io_wave(NonZeroUsize::new(2).unwrap())
959            .unwrap();
960
961        assert!(matches!(acquired, RuntimeIoWaveTryAcquire::Acquired(None)));
962    }
963}
hawdb_core/cancellation.rs

hawdb_core/
cancellation.rs
