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

hawdb_core/
cancellation.rs
