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::{
291    RuntimeMemoryController, RuntimeMemoryError, RuntimeMemoryPermit, RuntimeMemoryTarget,
292};
293
294#[derive(Debug, Clone)]
295pub struct RuntimeTaskContext {
296    cancellation: RuntimeCancellationToken,
297    deadline: Option<Instant>,
298    admitted_parallelism: NonZeroUsize,
299    executor_thread_limit: Option<NonZeroUsize>,
300    memory_reservation: Option<RuntimeMemoryReservation>,
301    memory_controller: Option<Arc<dyn RuntimeMemoryController>>,
302    memory_target: Option<Arc<dyn RuntimeMemoryTarget>>,
303    io_wave_controller: Option<Arc<dyn RuntimeIoWaveController>>,
304    work_budget: Option<Arc<work::RuntimeWorkBudget>>,
305}
306
307/// Memory already reserved for one task by the runtime governor.
308#[derive(Debug, Clone, Copy, PartialEq, Eq)]
309pub struct RuntimeMemoryReservation {
310    memory_bytes: u64,
311    result_bytes: u64,
312}
313
314impl RuntimeMemoryReservation {
315    pub const fn new(memory_bytes: u64, result_bytes: u64) -> Self {
316        Self {
317            memory_bytes,
318            result_bytes,
319        }
320    }
321
322    pub const fn memory_bytes(self) -> u64 {
323        self.memory_bytes
324    }
325
326    pub const fn result_bytes(self) -> u64 {
327        self.result_bytes
328    }
329
330    pub const fn intersect(self, other: Self) -> Self {
331        Self {
332            memory_bytes: if self.memory_bytes < other.memory_bytes {
333                self.memory_bytes
334            } else {
335                other.memory_bytes
336            },
337            result_bytes: if self.result_bytes < other.result_bytes {
338                self.result_bytes
339            } else {
340                other.result_bytes
341            },
342        }
343    }
344}
345
346impl RuntimeTaskContext {
347    pub fn new(cancellation: RuntimeCancellationToken, deadline: Option<Instant>) -> Self {
348        Self {
349            cancellation,
350            deadline,
351            admitted_parallelism: NonZeroUsize::MIN,
352            executor_thread_limit: None,
353            memory_reservation: None,
354            memory_controller: None,
355            memory_target: None,
356            io_wave_controller: None,
357            work_budget: None,
358        }
359    }
360
361    pub fn without_deadline(cancellation: RuntimeCancellationToken) -> Self {
362        Self::new(cancellation, None)
363    }
364
365    pub fn with_timeout(timeout: Duration) -> Self {
366        Self::new(
367            RuntimeCancellationToken::new(),
368            Instant::now().checked_add(timeout),
369        )
370    }
371
372    pub fn cancellation(&self) -> &RuntimeCancellationToken {
373        &self.cancellation
374    }
375
376    pub fn child(&self) -> Self {
377        Self {
378            cancellation: self.cancellation.child(),
379            deadline: self.deadline,
380            admitted_parallelism: self.admitted_parallelism,
381            executor_thread_limit: self.executor_thread_limit,
382            memory_reservation: self.memory_reservation,
383            memory_controller: self.memory_controller.clone(),
384            memory_target: self.memory_target.clone(),
385            io_wave_controller: self.io_wave_controller.clone(),
386            work_budget: self.work_budget.clone(),
387        }
388    }
389
390    /// Applies statement constraints without replacing an enclosing admission.
391    /// Cancellation from either scope stops the statement; cancelling the combined
392    /// scope does not cancel its parents. Budgets remain shared and can only tighten.
393    /// Independent controllers cannot represent one admission and are rejected.
394    pub fn constrained_by(&self, request: &Self) -> crate::Result<Self> {
395        fn controller<T: ?Sized>(
396            parent: &Option<Arc<T>>,
397            request: &Option<Arc<T>>,
398            resource: &str,
399        ) -> crate::Result<Option<Arc<T>>> {
400            match (parent, request) {
401                (Some(parent), Some(request)) if !Arc::ptr_eq(parent, request) => {
402                    Err(crate::HawDBError::Execution(format!(
403                        "request task context carries a different {resource} admission"
404                    )))
405                }
406                (Some(parent), _) => Ok(Some(parent.clone())),
407                (_, request) => Ok(request.clone()),
408            }
409        }
410        fn minimum<T: Ord + Copy>(left: Option<T>, right: Option<T>) -> Option<T> {
411            match (left, right) {
412                (Some(left), Some(right)) => Some(left.min(right)),
413                (left, right) => left.or(right),
414            }
415        }
416        Ok(Self {
417            cancellation: self.cancellation.linked(&request.cancellation),
418            deadline: minimum(self.deadline, request.deadline),
419            admitted_parallelism: self.admitted_parallelism.min(request.admitted_parallelism),
420            executor_thread_limit: minimum(
421                self.executor_thread_limit,
422                request.executor_thread_limit,
423            ),
424            memory_reservation: match (self.memory_reservation, request.memory_reservation) {
425                (Some(parent), Some(request)) => Some(parent.intersect(request)),
426                (parent, request) => parent.or(request),
427            },
428            memory_controller: controller(
429                &self.memory_controller,
430                &request.memory_controller,
431                "memory",
432            )?,
433            memory_target: controller(
434                &self.memory_target,
435                &request.memory_target,
436                "memory target",
437            )?,
438            io_wave_controller: controller(
439                &self.io_wave_controller,
440                &request.io_wave_controller,
441                "I/O",
442            )?,
443            work_budget: work::RuntimeWorkBudget::intersect(
444                &self.work_budget,
445                &request.work_budget,
446            ),
447        })
448    }
449
450    /// Carries the CPU parallelism already reserved by the runtime governor.
451    ///
452    /// This is an execution ceiling, not a request to create worker threads.
453    /// Executors must still apply their operator memory and input-size limits.
454    pub fn with_admitted_parallelism(mut self, admitted_parallelism: NonZeroUsize) -> Self {
455        self.admitted_parallelism = admitted_parallelism;
456        self
457    }
458
459    pub fn admitted_parallelism(&self) -> NonZeroUsize {
460        self.admitted_parallelism
461    }
462
463    /// Carries the physical executor-thread ceiling derived by the runtime governor.
464    ///
465    /// This is distinct from [`Self::admitted_parallelism`], which is the CPU
466    /// reservation for one task. An absent limit denotes an ungoverned library call.
467    pub fn with_executor_thread_limit(mut self, limit: NonZeroUsize) -> Self {
468        self.executor_thread_limit = Some(
469            self.executor_thread_limit
470                .map_or(limit, |current| current.min(limit)),
471        );
472        self
473    }
474
475    pub fn executor_thread_limit(&self) -> Option<NonZeroUsize> {
476        self.executor_thread_limit
477    }
478
479    /// Carries the memory already reserved by the runtime governor.
480    ///
481    /// Governed executors must use this reservation instead of a static
482    /// per-query configuration limit. An absent reservation denotes an
483    /// ungoverned library call and preserves the configured fallback.
484    pub fn with_memory_reservation(mut self, reservation: RuntimeMemoryReservation) -> Self {
485        self.memory_reservation = Some(reservation);
486        self
487    }
488
489    pub fn memory_reservation(&self) -> Option<RuntimeMemoryReservation> {
490        self.memory_reservation
491    }
492
493    /// Whether execution carries dynamic resource controllers rather than only
494    /// a static ceiling. Synchronous adapters must preserve these controllers
495    /// and must not replace an originating governor admission with a local cap.
496    pub fn has_resource_controllers(&self) -> bool {
497        self.memory_controller.is_some()
498            || self.memory_target.is_some()
499            || self.io_wave_controller.is_some()
500    }
501
502    /// Carries host-driven pressure feedback to revocable operator work areas.
503    pub fn with_memory_target(mut self, target: Arc<dyn RuntimeMemoryTarget>) -> Self {
504        self.memory_target = Some(target);
505        self
506    }
507
508    /// The target can only tighten the caller's ceiling. Operators retain a
509    /// bounded progress floor and continue to enforce all hard allocation limits.
510    pub fn working_memory_target(&self, hard_limit: u64) -> u64 {
511        self.memory_target.as_ref().map_or(hard_limit, |target| {
512            target.target_bytes(hard_limit).min(hard_limit)
513        })
514    }
515
516    /// Binds actual allocation ownership to an already admitted task.
517    pub fn with_memory_controller(mut self, controller: Arc<dyn RuntimeMemoryController>) -> Self {
518        self.memory_controller = Some(controller);
519        self
520    }
521
522    /// Reserve before allocating and retain the returned lease until destruction.
523    /// Clones and children share their controller. Ungoverned contexts return
524    /// `None`; a supplied static ceiling still limits an individual allocation.
525    pub fn reserve_working_memory(
526        &self,
527        bytes: u64,
528    ) -> Result<Option<Box<dyn RuntimeMemoryPermit>>, RuntimeMemoryError> {
529        self.checkpoint().map_err(RuntimeMemoryError::Stopped)?;
530        if let Some(reservation) = self.memory_reservation
531            && bytes > reservation.memory_bytes()
532        {
533            return Err(RuntimeMemoryError::ReservationExceeded {
534                requested_bytes: bytes,
535                available_bytes: reservation.memory_bytes(),
536            });
537        }
538        self.memory_controller
539            .as_ref()
540            .map(|controller| {
541                controller.reserve(
542                    bytes,
543                    self.memory_reservation
544                        .map_or(u64::MAX, RuntimeMemoryReservation::memory_bytes),
545                )
546            })
547            .transpose()
548    }
549
550    /// Binds the controller owned by a successful runtime admission.
551    pub fn with_io_wave_controller(mut self, controller: Arc<dyn RuntimeIoWaveController>) -> Self {
552        self.io_wave_controller = Some(controller);
553        self
554    }
555
556    /// Acquires capacity for one storage I/O wave when this is an admitted context.
557    ///
558    /// Raw contexts have no controller and return `Ok(None)`. The returned
559    /// permit must remain live only while physical reads are in flight.
560    pub fn acquire_io_wave(
561        &self,
562        slots: NonZeroUsize,
563    ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
564        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
565        self.io_wave_controller
566            .as_ref()
567            .map(|controller| controller.acquire(slots, self))
568            .transpose()
569    }
570
571    /// Attempts to acquire capacity for one storage I/O wave without blocking.
572    ///
573    /// Raw, ungoverned contexts return `Acquired(None)`. Governed contexts
574    /// return `Pending` while another wave owns the requested capacity.
575    pub fn try_acquire_io_wave(
576        &self,
577        slots: NonZeroUsize,
578    ) -> Result<RuntimeIoWaveTryAcquire, RuntimeIoWaveError> {
579        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
580        let Some(controller) = &self.io_wave_controller else {
581            return Ok(RuntimeIoWaveTryAcquire::Acquired(None));
582        };
583        controller.try_acquire(slots, self).map(|permit| {
584            permit.map_or(RuntimeIoWaveTryAcquire::Pending, |permit| {
585                RuntimeIoWaveTryAcquire::Acquired(Some(permit))
586            })
587        })
588    }
589
590    /// Starts a bounded execution scope while retaining every parent limit and
591    /// the original cancellation, memory and I/O admission. Clones share usage.
592    pub fn with_work_limit(mut self, limit: usize) -> Self {
593        self.work_budget = Some(Arc::new(work::RuntimeWorkBudget::new(
594            limit,
595            self.work_budget.take(),
596        )));
597        self
598    }
599
600    pub fn has_work_limit(&self) -> bool {
601        self.work_budget.is_some()
602    }
603
604    pub fn work_units(&self) -> usize {
605        self.work_budget.as_ref().map_or(0, |budget| budget.used())
606    }
607
608    /// Preserve a typed refusal across adapters whose error transport is text.
609    pub fn check_work_limit(&self) -> crate::Result<()> {
610        match self
611            .work_budget
612            .as_ref()
613            .and_then(|budget| budget.refusal())
614        {
615            Some(error) => Err(error.into()),
616            None => Ok(()),
617        }
618    }
619
620    /// Charge explicit candidate work before examining it, including rejected
621    /// candidates. Checkpoints and allocation byte counts are not work units.
622    pub fn charge_work(&self, units: usize) -> crate::Result<()> {
623        let Some(budget) = &self.work_budget else {
624            return Ok(());
625        };
626        self.checkpoint().map_err(|reason| {
627            crate::HawDBError::Execution(format!("runtime task stopped: {reason}"))
628        })?;
629        budget.charge(units)
630    }
631
632    pub fn deadline(&self) -> Option<Instant> {
633        self.deadline
634    }
635
636    pub fn remaining(&self) -> Option<Duration> {
637        self.deadline
638            .map(|deadline| deadline.saturating_duration_since(Instant::now()))
639    }
640
641    pub fn checkpoint(&self) -> Result<(), RuntimeCancellationReason> {
642        if let Some(reason) = self.cancellation.reason() {
643            return Err(reason);
644        }
645        if self
646            .deadline
647            .is_some_and(|deadline| Instant::now() >= deadline)
648        {
649            return Err(RuntimeCancellationReason::DeadlineExceeded);
650        }
651        Ok(())
652    }
653}
654
655impl Default for RuntimeTaskContext {
656    fn default() -> Self {
657        Self::without_deadline(RuntimeCancellationToken::new())
658    }
659}
660
661#[derive(Debug, Clone, Copy, PartialEq, Eq)]
662pub enum RuntimeCancellationReason {
663    Cancelled,
664    DeadlineExceeded,
665    ResourceExhausted,
666}
667
668impl RuntimeCancellationReason {
669    pub const fn as_str(self) -> &'static str {
670        match self {
671            Self::Cancelled => "cancelled",
672            Self::DeadlineExceeded => "deadline_exceeded",
673            Self::ResourceExhausted => "resource_exhausted",
674        }
675    }
676}
677
678impl Display for RuntimeCancellationReason {
679    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
680        formatter.write_str(self.as_str())
681    }
682}
683
684impl Error for RuntimeCancellationReason {}
685
686#[cfg(test)]
687mod tests {
688    use super::*;
689    use std::sync::atomic::AtomicUsize;
690    use std::sync::atomic::Ordering;
691    use std::task::Wake;
692
693    #[derive(Debug)]
694    struct RecordingIoController {
695        active_slots: Arc<AtomicUsize>,
696    }
697
698    #[derive(Debug)]
699    struct RecordingIoPermit {
700        active_slots: Arc<AtomicUsize>,
701        slots: usize,
702    }
703
704    impl Drop for RecordingIoPermit {
705        fn drop(&mut self) {
706            self.active_slots.fetch_sub(self.slots, Ordering::AcqRel);
707        }
708    }
709
710    impl RuntimeIoWaveController for RecordingIoController {
711        fn try_acquire(
712            &self,
713            slots: NonZeroUsize,
714            context: &RuntimeTaskContext,
715        ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
716            self.acquire(slots, context).map(Some)
717        }
718
719        fn acquire(
720            &self,
721            slots: NonZeroUsize,
722            context: &RuntimeTaskContext,
723        ) -> Result<Box<dyn RuntimeIoWavePermit>, RuntimeIoWaveError> {
724            context.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
725            self.active_slots.fetch_add(slots.get(), Ordering::AcqRel);
726            Ok(Box::new(RecordingIoPermit {
727                active_slots: Arc::clone(&self.active_slots),
728                slots: slots.get(),
729            }))
730        }
731    }
732
733    #[test]
734    fn cancellation_is_shared_across_context_clones() {
735        let token = RuntimeCancellationToken::new();
736        let context = RuntimeTaskContext::without_deadline(token.clone());
737        assert!(context.checkpoint().is_ok());
738        assert!(token.cancel());
739        assert!(!token.cancel());
740        assert_eq!(
741            context.checkpoint(),
742            Err(RuntimeCancellationReason::Cancelled)
743        );
744    }
745
746    #[test]
747    fn expired_deadline_fails_at_a_cooperative_checkpoint() {
748        let context = RuntimeTaskContext::new(
749            RuntimeCancellationToken::new(),
750            Some(Instant::now() - Duration::from_millis(1)),
751        );
752        assert_eq!(
753            context.checkpoint(),
754            Err(RuntimeCancellationReason::DeadlineExceeded)
755        );
756    }
757
758    #[test]
759    fn child_cancellation_is_local_and_parent_cancellation_propagates() {
760        let parent_token = RuntimeCancellationToken::new();
761        let parent = RuntimeTaskContext::without_deadline(parent_token.clone());
762        let child = parent.child();
763
764        assert!(child.cancellation().cancel());
765        assert!(child.checkpoint().is_err());
766        assert!(parent.checkpoint().is_ok());
767
768        let sibling = parent.child();
769        assert!(parent_token.cancel());
770        assert!(parent.checkpoint().is_err());
771        assert!(sibling.checkpoint().is_err());
772    }
773
774    #[test]
775    fn resource_exhaustion_keeps_its_cause_and_is_local_to_the_child() {
776        let parent = RuntimeTaskContext::default();
777        let child = parent.child();
778        let clone = child.clone();
779        let grandchild = child.child();
780        assert!(child.cancellation().cancel_for_resource_exhaustion());
781        assert!(!clone.cancellation().cancel());
782        assert!(!clone.cancellation().cancel_for_resource_exhaustion());
783        for context in [&child, &clone, &grandchild] {
784            assert_eq!(
785                context.checkpoint(),
786                Err(RuntimeCancellationReason::ResourceExhausted)
787            );
788        }
789        assert_eq!(parent.checkpoint(), Ok(()));
790        assert_eq!(parent.child().checkpoint(), Ok(()));
791
792        let cancelled = parent.child();
793        assert!(cancelled.cancellation().cancel());
794        assert!(!cancelled.cancellation().cancel_for_resource_exhaustion());
795        assert_eq!(
796            cancelled.checkpoint(),
797            Err(RuntimeCancellationReason::Cancelled)
798        );
799    }
800
801    #[derive(Default)]
802    struct WakeCounter(AtomicUsize);
803
804    impl Wake for WakeCounter {
805        fn wake(self: Arc<Self>) {
806            self.0.fetch_add(1, Ordering::AcqRel);
807        }
808    }
809
810    #[test]
811    fn constrained_context_observes_both_cancellations_without_cancelling_parents() {
812        for cancel_parent in [false, true] {
813            let parent = RuntimeTaskContext::default();
814            let request = RuntimeTaskContext::default();
815            let combined = parent.constrained_by(&request).unwrap();
816            let counter = Arc::new(WakeCounter::default());
817            let waker = Waker::from(counter.clone());
818            let mut context = Context::from_waker(&waker);
819            let mut cancelled = Box::pin(combined.cancellation().cancelled());
820            assert!(cancelled.as_mut().poll(&mut context).is_pending());
821            if cancel_parent {
822                parent.cancellation().cancel();
823            } else {
824                request.cancellation().cancel();
825            }
826            assert_eq!(counter.0.load(Ordering::Acquire), 1);
827            assert!(cancelled.as_mut().poll(&mut context).is_ready());
828            assert_eq!(
829                combined.checkpoint(),
830                Err(RuntimeCancellationReason::Cancelled)
831            );
832        }
833        let parent = RuntimeTaskContext::default();
834        let request = RuntimeTaskContext::default();
835        let combined = parent.constrained_by(&request).unwrap();
836        combined.cancellation().cancel();
837        assert!(parent.checkpoint().is_ok());
838        assert!(request.checkpoint().is_ok());
839    }
840
841    #[test]
842    fn constrained_descendants_register_shared_ancestors_once_and_unregister_on_drop() {
843        let parent = RuntimeTaskContext::default();
844        let request = parent.child();
845        let combined = parent.constrained_by(&request).unwrap();
846        let waker = Waker::from(Arc::new(WakeCounter::default()));
847        let mut context = Context::from_waker(&waker);
848        let mut cancelled = Box::pin(combined.cancellation().cancelled());
849        assert!(cancelled.as_mut().poll(&mut context).is_pending());
850        assert_eq!(parent.cancellation.state.waiters.lock().unwrap().len(), 1);
851        assert_eq!(request.cancellation.state.waiters.lock().unwrap().len(), 1);
852        drop(cancelled);
853        assert!(parent.cancellation.state.waiters.lock().unwrap().is_empty());
854        assert!(request
855            .cancellation
856            .state
857            .waiters
858            .lock()
859            .unwrap()
860            .is_empty());
861    }
862
863    #[test]
864    fn constrained_context_tightens_limits_and_keeps_the_original_io_controller() {
865        let active_slots = Arc::new(AtomicUsize::new(0));
866        let parent = RuntimeTaskContext::with_timeout(Duration::from_secs(5))
867            .with_memory_reservation(RuntimeMemoryReservation::new(100, 20))
868            .with_admitted_parallelism(NonZeroUsize::new(4).unwrap())
869            .with_io_wave_controller(Arc::new(RecordingIoController {
870                active_slots: active_slots.clone(),
871            }));
872        let request = RuntimeTaskContext::with_timeout(Duration::from_secs(10))
873            .with_memory_reservation(RuntimeMemoryReservation::new(30, 50))
874            .with_admitted_parallelism(NonZeroUsize::new(2).unwrap());
875        let combined = parent.constrained_by(&request).unwrap();
876        assert_eq!(combined.deadline(), parent.deadline());
877        assert_eq!(
878            combined.memory_reservation(),
879            Some(RuntimeMemoryReservation::new(30, 20))
880        );
881        assert_eq!(combined.admitted_parallelism().get(), 2);
882        let io = combined.acquire_io_wave(NonZeroUsize::MIN).unwrap();
883        assert_eq!(active_slots.load(Ordering::Acquire), 1);
884        drop(io);
885        assert_eq!(active_slots.load(Ordering::Acquire), 0);
886        assert!(parent.constrained_by(&parent.child()).is_ok());
887        let independent =
888            request.with_io_wave_controller(Arc::new(RecordingIoController { active_slots }));
889        assert!(parent.constrained_by(&independent).is_err());
890    }
891
892    #[test]
893    fn cancellation_future_is_woken_by_parent_cancellation() {
894        let parent = RuntimeCancellationToken::new();
895        let child = parent.child();
896        let counter = Arc::new(WakeCounter::default());
897        let waker = Waker::from(Arc::clone(&counter));
898        let mut context = Context::from_waker(&waker);
899        let mut cancelled = Box::pin(child.cancelled());
900
901        assert!(cancelled.as_mut().poll(&mut context).is_pending());
902        assert!(parent.cancel());
903        assert_eq!(counter.0.load(Ordering::Acquire), 1);
904        assert!(cancelled.as_mut().poll(&mut context).is_ready());
905    }
906
907    #[test]
908    fn resource_exhaustion_wakes_waiters_on_descendants() {
909        let parent = RuntimeTaskContext::default();
910        let child = parent.child();
911        let counter = Arc::new(WakeCounter::default());
912        let waker = Waker::from(Arc::clone(&counter));
913        let mut context = Context::from_waker(&waker);
914        let mut cancelled = Box::pin(child.cancellation().cancelled());
915        assert!(cancelled.as_mut().poll(&mut context).is_pending());
916        assert!(parent.cancellation().cancel_for_resource_exhaustion());
917        assert_eq!(counter.0.load(Ordering::Acquire), 1);
918        assert!(cancelled.as_mut().poll(&mut context).is_ready());
919        assert_eq!(
920            child.checkpoint(),
921            Err(RuntimeCancellationReason::ResourceExhausted)
922        );
923    }
924
925    #[test]
926    fn dropped_cancellation_future_unregisters_from_its_parent_chain() {
927        let parent = RuntimeCancellationToken::new();
928        let child = parent.child();
929        let counter = Arc::new(WakeCounter::default());
930        let waker = Waker::from(Arc::clone(&counter));
931        let mut context = Context::from_waker(&waker);
932        let mut cancelled = Box::pin(child.cancelled());
933
934        assert!(cancelled.as_mut().poll(&mut context).is_pending());
935        drop(cancelled);
936        assert!(parent.cancel());
937        assert_eq!(counter.0.load(Ordering::Acquire), 0);
938    }
939
940    #[test]
941    fn child_preserves_admitted_parallelism() {
942        let context = RuntimeTaskContext::default()
943            .with_admitted_parallelism(NonZeroUsize::new(4).expect("test parallelism is non-zero"))
944            .with_executor_thread_limit(
945                NonZeroUsize::new(2).expect("test thread limit is non-zero"),
946            );
947
948        assert_eq!(context.admitted_parallelism().get(), 4);
949        assert_eq!(context.child().admitted_parallelism().get(), 4);
950        assert_eq!(context.child().executor_thread_limit().unwrap().get(), 2);
951    }
952
953    #[test]
954    fn child_preserves_memory_reservation() {
955        let reservation = RuntimeMemoryReservation::new(128, 32);
956        let context = RuntimeTaskContext::default().with_memory_reservation(reservation);
957
958        assert_eq!(context.memory_reservation(), Some(reservation));
959        assert_eq!(context.child().memory_reservation(), Some(reservation));
960    }
961
962    #[test]
963    fn child_preserves_io_controller_and_permit_releases_slots() {
964        let active_slots = Arc::new(AtomicUsize::new(0));
965        let context = RuntimeTaskContext::default().with_io_wave_controller(Arc::new(
966            RecordingIoController {
967                active_slots: Arc::clone(&active_slots),
968            },
969        ));
970
971        let permit = context
972            .child()
973            .acquire_io_wave(NonZeroUsize::new(2).unwrap())
974            .unwrap()
975            .unwrap();
976        assert_eq!(active_slots.load(Ordering::Acquire), 2);
977        drop(permit);
978        assert_eq!(active_slots.load(Ordering::Acquire), 0);
979    }
980
981    #[test]
982    fn ungoverned_try_acquire_is_immediately_ready() {
983        let acquired = RuntimeTaskContext::default()
984            .try_acquire_io_wave(NonZeroUsize::new(2).unwrap())
985            .unwrap();
986
987        assert!(matches!(acquired, RuntimeIoWaveTryAcquire::Acquired(None)));
988    }
989}
hawdb_core/cancellation.rs

hawdb_core/
cancellation.rs
