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;
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    parent: Option<RuntimeCancellationToken>,
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                parent: 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                .parent
107                .as_ref()
108                .and_then(RuntimeCancellationToken::reason),
109        }
110    }
111
112    pub fn child(&self) -> Self {
113        Self {
114            state: Arc::new(RuntimeCancellationState {
115                cancellation_reason: AtomicU8::new(NOT_CANCELLED),
116                parent: Some(self.clone()),
117                next_waiter_id: AtomicU64::new(1),
118                waiters: Mutex::new(BTreeMap::new()),
119            }),
120        }
121    }
122
123    /// Resolves when this token or any parent token is cancelled.
124    pub fn cancelled(&self) -> RuntimeCancellationFuture<'_> {
125        RuntimeCancellationFuture {
126            token: self,
127            registrations: Vec::new(),
128        }
129    }
130}
131
132struct RuntimeCancellationRegistration {
133    state: Weak<RuntimeCancellationState>,
134    waiter_id: u64,
135}
136
137/// A runtime-neutral future returned by [`RuntimeCancellationToken::cancelled`].
138pub struct RuntimeCancellationFuture<'a> {
139    token: &'a RuntimeCancellationToken,
140    registrations: Vec<RuntimeCancellationRegistration>,
141}
142
143impl Future for RuntimeCancellationFuture<'_> {
144    type Output = ();
145
146    fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
147        if self.token.is_cancelled() {
148            return Poll::Ready(());
149        }
150
151        if self.registrations.is_empty() {
152            let mut token = Some(self.token);
153            while let Some(current) = token {
154                let waiter_id = current.state.next_waiter_id.fetch_add(1, Ordering::Relaxed);
155                current
156                    .state
157                    .waiters
158                    .lock()
159                    .unwrap_or_else(|poisoned| poisoned.into_inner())
160                    .insert(waiter_id, context.waker().clone());
161                self.registrations.push(RuntimeCancellationRegistration {
162                    state: Arc::downgrade(&current.state),
163                    waiter_id,
164                });
165                token = current.state.parent.as_ref();
166            }
167        } else {
168            for registration in &self.registrations {
169                if let Some(state) = registration.state.upgrade()
170                    && let Some(waiter) = state
171                        .waiters
172                        .lock()
173                        .unwrap_or_else(|poisoned| poisoned.into_inner())
174                        .get_mut(&registration.waiter_id)
175                    && !waiter.will_wake(context.waker())
176                {
177                    *waiter = context.waker().clone();
178                }
179            }
180        }
181
182        if self.token.is_cancelled() {
183            Poll::Ready(())
184        } else {
185            Poll::Pending
186        }
187    }
188}
189
190impl Drop for RuntimeCancellationFuture<'_> {
191    fn drop(&mut self) {
192        for registration in &self.registrations {
193            if let Some(state) = registration.state.upgrade() {
194                state
195                    .waiters
196                    .lock()
197                    .unwrap_or_else(|poisoned| poisoned.into_inner())
198                    .remove(&registration.waiter_id);
199            }
200        }
201    }
202}
203
204/// A drop guard for one admitted storage I/O wave.
205pub trait RuntimeIoWavePermit: Debug + Send {}
206
207impl<T: Debug + Send> RuntimeIoWavePermit for T {}
208
209/// Execution-owned hook for acquiring the I/O capacity declared at admission.
210pub trait RuntimeIoWaveController: Debug + Send + Sync {
211    /// Attempts to acquire one I/O wave without blocking the calling thread.
212    ///
213    /// `Ok(None)` means that the reservation is currently saturated. Async
214    /// runtime adapters must yield before retrying instead of spinning.
215    fn try_acquire(
216        &self,
217        slots: NonZeroUsize,
218        context: &RuntimeTaskContext,
219    ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError>;
220
221    fn acquire(
222        &self,
223        slots: NonZeroUsize,
224        context: &RuntimeTaskContext,
225    ) -> Result<Box<dyn RuntimeIoWavePermit>, RuntimeIoWaveError>;
226}
227
228#[derive(Debug, Clone, Copy, PartialEq, Eq)]
229pub enum RuntimeIoWaveError {
230    Stopped(RuntimeCancellationReason),
231    ReservationExceeded {
232        requested_slots: usize,
233        reserved_slots: usize,
234    },
235}
236
237#[derive(Debug)]
238pub enum RuntimeIoWaveTryAcquire {
239    /// Capacity is available. Ungoverned contexts carry no permit.
240    Acquired(Option<Box<dyn RuntimeIoWavePermit>>),
241    /// Capacity is currently saturated and the caller should yield.
242    Pending,
243}
244
245impl Display for RuntimeIoWaveError {
246    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
247        match self {
248            Self::Stopped(reason) => write!(formatter, "runtime I/O wave stopped: {reason}"),
249            Self::ReservationExceeded {
250                requested_slots,
251                reserved_slots,
252            } => write!(
253                formatter,
254                "runtime I/O wave requested {requested_slots} slots, exceeding the {reserved_slots}-slot reservation"
255            ),
256        }
257    }
258}
259
260impl Error for RuntimeIoWaveError {
261    fn source(&self) -> Option<&(dyn Error + 'static)> {
262        match self {
263            Self::Stopped(reason) => Some(reason),
264            Self::ReservationExceeded { .. } => None,
265        }
266    }
267}
268
269mod memory;
270mod work;
271pub use memory::{RuntimeMemoryController, RuntimeMemoryError, RuntimeMemoryPermit};
272
273#[derive(Debug, Clone)]
274pub struct RuntimeTaskContext {
275    cancellation: RuntimeCancellationToken,
276    deadline: Option<Instant>,
277    admitted_parallelism: NonZeroUsize,
278    executor_thread_limit: Option<NonZeroUsize>,
279    memory_reservation: Option<RuntimeMemoryReservation>,
280    memory_controller: Option<Arc<dyn RuntimeMemoryController>>,
281    io_wave_controller: Option<Arc<dyn RuntimeIoWaveController>>,
282    work_budget: Option<Arc<work::RuntimeWorkBudget>>,
283}
284
285/// Memory already reserved for one task by the runtime governor.
286#[derive(Debug, Clone, Copy, PartialEq, Eq)]
287pub struct RuntimeMemoryReservation {
288    memory_bytes: u64,
289    result_bytes: u64,
290}
291
292impl RuntimeMemoryReservation {
293    pub const fn new(memory_bytes: u64, result_bytes: u64) -> Self {
294        Self {
295            memory_bytes,
296            result_bytes,
297        }
298    }
299
300    pub const fn memory_bytes(self) -> u64 {
301        self.memory_bytes
302    }
303
304    pub const fn result_bytes(self) -> u64 {
305        self.result_bytes
306    }
307
308    pub const fn intersect(self, other: Self) -> Self {
309        Self {
310            memory_bytes: if self.memory_bytes < other.memory_bytes {
311                self.memory_bytes
312            } else {
313                other.memory_bytes
314            },
315            result_bytes: if self.result_bytes < other.result_bytes {
316                self.result_bytes
317            } else {
318                other.result_bytes
319            },
320        }
321    }
322}
323
324impl RuntimeTaskContext {
325    pub fn new(cancellation: RuntimeCancellationToken, deadline: Option<Instant>) -> Self {
326        Self {
327            cancellation,
328            deadline,
329            admitted_parallelism: NonZeroUsize::MIN,
330            executor_thread_limit: None,
331            memory_reservation: None,
332            memory_controller: None,
333            io_wave_controller: None,
334            work_budget: None,
335        }
336    }
337
338    pub fn without_deadline(cancellation: RuntimeCancellationToken) -> Self {
339        Self::new(cancellation, None)
340    }
341
342    pub fn with_timeout(timeout: Duration) -> Self {
343        Self::new(
344            RuntimeCancellationToken::new(),
345            Instant::now().checked_add(timeout),
346        )
347    }
348
349    pub fn cancellation(&self) -> &RuntimeCancellationToken {
350        &self.cancellation
351    }
352
353    pub fn child(&self) -> Self {
354        Self {
355            cancellation: self.cancellation.child(),
356            deadline: self.deadline,
357            admitted_parallelism: self.admitted_parallelism,
358            executor_thread_limit: self.executor_thread_limit,
359            memory_reservation: self.memory_reservation,
360            memory_controller: self.memory_controller.clone(),
361            io_wave_controller: self.io_wave_controller.clone(),
362            work_budget: self.work_budget.clone(),
363        }
364    }
365
366    /// Carries the CPU parallelism already reserved by the runtime governor.
367    ///
368    /// This is an execution ceiling, not a request to create worker threads.
369    /// Executors must still apply their operator memory and input-size limits.
370    pub fn with_admitted_parallelism(mut self, admitted_parallelism: NonZeroUsize) -> Self {
371        self.admitted_parallelism = admitted_parallelism;
372        self
373    }
374
375    pub fn admitted_parallelism(&self) -> NonZeroUsize {
376        self.admitted_parallelism
377    }
378
379    /// Carries the physical executor-thread ceiling derived by the runtime governor.
380    ///
381    /// This is distinct from [`Self::admitted_parallelism`], which is the CPU
382    /// reservation for one task. An absent limit denotes an ungoverned library call.
383    pub fn with_executor_thread_limit(mut self, limit: NonZeroUsize) -> Self {
384        self.executor_thread_limit = Some(
385            self.executor_thread_limit
386                .map_or(limit, |current| current.min(limit)),
387        );
388        self
389    }
390
391    pub fn executor_thread_limit(&self) -> Option<NonZeroUsize> {
392        self.executor_thread_limit
393    }
394
395    /// Carries the memory already reserved by the runtime governor.
396    ///
397    /// Governed executors must use this reservation instead of a static
398    /// per-query configuration limit. An absent reservation denotes an
399    /// ungoverned library call and preserves the configured fallback.
400    pub fn with_memory_reservation(mut self, reservation: RuntimeMemoryReservation) -> Self {
401        self.memory_reservation = Some(reservation);
402        self
403    }
404
405    pub fn memory_reservation(&self) -> Option<RuntimeMemoryReservation> {
406        self.memory_reservation
407    }
408
409    /// Whether execution carries dynamic resource controllers rather than only
410    /// a static ceiling. Synchronous adapters must preserve these controllers
411    /// and must not replace an originating governor admission with a local cap.
412    pub fn has_resource_controllers(&self) -> bool {
413        self.memory_controller.is_some() || self.io_wave_controller.is_some()
414    }
415
416    /// Binds actual allocation ownership to an already admitted task.
417    pub fn with_memory_controller(mut self, controller: Arc<dyn RuntimeMemoryController>) -> Self {
418        self.memory_controller = Some(controller);
419        self
420    }
421
422    /// Reserve before allocating and retain the returned lease until destruction.
423    /// Clones and children share their controller. Ungoverned contexts return
424    /// `None`; a supplied static ceiling still limits an individual allocation.
425    pub fn reserve_working_memory(
426        &self,
427        bytes: u64,
428    ) -> Result<Option<Box<dyn RuntimeMemoryPermit>>, RuntimeMemoryError> {
429        self.checkpoint().map_err(RuntimeMemoryError::Stopped)?;
430        if let Some(reservation) = self.memory_reservation
431            && bytes > reservation.memory_bytes()
432        {
433            return Err(RuntimeMemoryError::ReservationExceeded {
434                requested_bytes: bytes,
435                available_bytes: reservation.memory_bytes(),
436            });
437        }
438        self.memory_controller
439            .as_ref()
440            .map(|controller| {
441                controller.reserve(
442                    bytes,
443                    self.memory_reservation
444                        .map_or(u64::MAX, RuntimeMemoryReservation::memory_bytes),
445                )
446            })
447            .transpose()
448    }
449
450    /// Binds the controller owned by a successful runtime admission.
451    pub fn with_io_wave_controller(mut self, controller: Arc<dyn RuntimeIoWaveController>) -> Self {
452        self.io_wave_controller = Some(controller);
453        self
454    }
455
456    /// Acquires capacity for one storage I/O wave when this is an admitted context.
457    ///
458    /// Raw contexts have no controller and return `Ok(None)`. The returned
459    /// permit must remain live only while physical reads are in flight.
460    pub fn acquire_io_wave(
461        &self,
462        slots: NonZeroUsize,
463    ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
464        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
465        self.io_wave_controller
466            .as_ref()
467            .map(|controller| controller.acquire(slots, self))
468            .transpose()
469    }
470
471    /// Attempts to acquire capacity for one storage I/O wave without blocking.
472    ///
473    /// Raw, ungoverned contexts return `Acquired(None)`. Governed contexts
474    /// return `Pending` while another wave owns the requested capacity.
475    pub fn try_acquire_io_wave(
476        &self,
477        slots: NonZeroUsize,
478    ) -> Result<RuntimeIoWaveTryAcquire, RuntimeIoWaveError> {
479        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
480        let Some(controller) = &self.io_wave_controller else {
481            return Ok(RuntimeIoWaveTryAcquire::Acquired(None));
482        };
483        controller.try_acquire(slots, self).map(|permit| {
484            permit.map_or(RuntimeIoWaveTryAcquire::Pending, |permit| {
485                RuntimeIoWaveTryAcquire::Acquired(Some(permit))
486            })
487        })
488    }
489
490    /// Starts a bounded execution scope while retaining every parent limit and
491    /// the original cancellation, memory and I/O admission. Clones share usage.
492    pub fn with_work_limit(mut self, limit: usize) -> Self {
493        self.work_budget = Some(Arc::new(work::RuntimeWorkBudget::new(
494            limit,
495            self.work_budget.take(),
496        )));
497        self
498    }
499
500    pub fn has_work_limit(&self) -> bool {
501        self.work_budget.is_some()
502    }
503
504    pub fn work_units(&self) -> usize {
505        self.work_budget.as_ref().map_or(0, |budget| budget.used())
506    }
507
508    /// Preserve a typed refusal across adapters whose error transport is text.
509    pub fn check_work_limit(&self) -> crate::Result<()> {
510        match self
511            .work_budget
512            .as_ref()
513            .and_then(|budget| budget.refusal())
514        {
515            Some(error) => Err(error.into()),
516            None => Ok(()),
517        }
518    }
519
520    /// Charge explicit candidate work before examining it, including rejected
521    /// candidates. Checkpoints and allocation byte counts are not work units.
522    pub fn charge_work(&self, units: usize) -> crate::Result<()> {
523        let Some(budget) = &self.work_budget else {
524            return Ok(());
525        };
526        self.checkpoint().map_err(|reason| {
527            crate::HawDBError::Execution(format!("runtime task stopped: {reason}"))
528        })?;
529        budget.charge(units)
530    }
531
532    pub fn deadline(&self) -> Option<Instant> {
533        self.deadline
534    }
535
536    pub fn remaining(&self) -> Option<Duration> {
537        self.deadline
538            .map(|deadline| deadline.saturating_duration_since(Instant::now()))
539    }
540
541    pub fn checkpoint(&self) -> Result<(), RuntimeCancellationReason> {
542        if let Some(reason) = self.cancellation.reason() {
543            return Err(reason);
544        }
545        if self
546            .deadline
547            .is_some_and(|deadline| Instant::now() >= deadline)
548        {
549            return Err(RuntimeCancellationReason::DeadlineExceeded);
550        }
551        Ok(())
552    }
553}
554
555impl Default for RuntimeTaskContext {
556    fn default() -> Self {
557        Self::without_deadline(RuntimeCancellationToken::new())
558    }
559}
560
561#[derive(Debug, Clone, Copy, PartialEq, Eq)]
562pub enum RuntimeCancellationReason {
563    Cancelled,
564    DeadlineExceeded,
565    ResourceExhausted,
566}
567
568impl RuntimeCancellationReason {
569    pub const fn as_str(self) -> &'static str {
570        match self {
571            Self::Cancelled => "cancelled",
572            Self::DeadlineExceeded => "deadline_exceeded",
573            Self::ResourceExhausted => "resource_exhausted",
574        }
575    }
576}
577
578impl Display for RuntimeCancellationReason {
579    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
580        formatter.write_str(self.as_str())
581    }
582}
583
584impl Error for RuntimeCancellationReason {}
585
586#[cfg(test)]
587mod tests {
588    use super::*;
589    use std::sync::atomic::AtomicUsize;
590    use std::sync::atomic::Ordering;
591    use std::task::Wake;
592
593    #[derive(Debug)]
594    struct RecordingIoController {
595        active_slots: Arc<AtomicUsize>,
596    }
597
598    #[derive(Debug)]
599    struct RecordingIoPermit {
600        active_slots: Arc<AtomicUsize>,
601        slots: usize,
602    }
603
604    impl Drop for RecordingIoPermit {
605        fn drop(&mut self) {
606            self.active_slots.fetch_sub(self.slots, Ordering::AcqRel);
607        }
608    }
609
610    impl RuntimeIoWaveController for RecordingIoController {
611        fn try_acquire(
612            &self,
613            slots: NonZeroUsize,
614            context: &RuntimeTaskContext,
615        ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
616            self.acquire(slots, context).map(Some)
617        }
618
619        fn acquire(
620            &self,
621            slots: NonZeroUsize,
622            context: &RuntimeTaskContext,
623        ) -> Result<Box<dyn RuntimeIoWavePermit>, RuntimeIoWaveError> {
624            context.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
625            self.active_slots.fetch_add(slots.get(), Ordering::AcqRel);
626            Ok(Box::new(RecordingIoPermit {
627                active_slots: Arc::clone(&self.active_slots),
628                slots: slots.get(),
629            }))
630        }
631    }
632
633    #[test]
634    fn cancellation_is_shared_across_context_clones() {
635        let token = RuntimeCancellationToken::new();
636        let context = RuntimeTaskContext::without_deadline(token.clone());
637        assert!(context.checkpoint().is_ok());
638        assert!(token.cancel());
639        assert!(!token.cancel());
640        assert_eq!(
641            context.checkpoint(),
642            Err(RuntimeCancellationReason::Cancelled)
643        );
644    }
645
646    #[test]
647    fn expired_deadline_fails_at_a_cooperative_checkpoint() {
648        let context = RuntimeTaskContext::new(
649            RuntimeCancellationToken::new(),
650            Some(Instant::now() - Duration::from_millis(1)),
651        );
652        assert_eq!(
653            context.checkpoint(),
654            Err(RuntimeCancellationReason::DeadlineExceeded)
655        );
656    }
657
658    #[test]
659    fn child_cancellation_is_local_and_parent_cancellation_propagates() {
660        let parent_token = RuntimeCancellationToken::new();
661        let parent = RuntimeTaskContext::without_deadline(parent_token.clone());
662        let child = parent.child();
663
664        assert!(child.cancellation().cancel());
665        assert!(child.checkpoint().is_err());
666        assert!(parent.checkpoint().is_ok());
667
668        let sibling = parent.child();
669        assert!(parent_token.cancel());
670        assert!(parent.checkpoint().is_err());
671        assert!(sibling.checkpoint().is_err());
672    }
673
674    #[test]
675    fn resource_exhaustion_keeps_its_cause_and_is_local_to_the_child() {
676        let parent = RuntimeTaskContext::default();
677        let child = parent.child();
678        let clone = child.clone();
679        let grandchild = child.child();
680        assert!(child.cancellation().cancel_for_resource_exhaustion());
681        assert!(!clone.cancellation().cancel());
682        assert!(!clone.cancellation().cancel_for_resource_exhaustion());
683        for context in [&child, &clone, &grandchild] {
684            assert_eq!(
685                context.checkpoint(),
686                Err(RuntimeCancellationReason::ResourceExhausted)
687            );
688        }
689        assert_eq!(parent.checkpoint(), Ok(()));
690        assert_eq!(parent.child().checkpoint(), Ok(()));
691
692        let cancelled = parent.child();
693        assert!(cancelled.cancellation().cancel());
694        assert!(!cancelled.cancellation().cancel_for_resource_exhaustion());
695        assert_eq!(
696            cancelled.checkpoint(),
697            Err(RuntimeCancellationReason::Cancelled)
698        );
699    }
700
701    #[derive(Default)]
702    struct WakeCounter(AtomicUsize);
703
704    impl Wake for WakeCounter {
705        fn wake(self: Arc<Self>) {
706            self.0.fetch_add(1, Ordering::AcqRel);
707        }
708    }
709
710    #[test]
711    fn cancellation_future_is_woken_by_parent_cancellation() {
712        let parent = RuntimeCancellationToken::new();
713        let child = parent.child();
714        let counter = Arc::new(WakeCounter::default());
715        let waker = Waker::from(Arc::clone(&counter));
716        let mut context = Context::from_waker(&waker);
717        let mut cancelled = Box::pin(child.cancelled());
718
719        assert!(cancelled.as_mut().poll(&mut context).is_pending());
720        assert!(parent.cancel());
721        assert_eq!(counter.0.load(Ordering::Acquire), 1);
722        assert!(cancelled.as_mut().poll(&mut context).is_ready());
723    }
724
725    #[test]
726    fn resource_exhaustion_wakes_waiters_on_descendants() {
727        let parent = RuntimeTaskContext::default();
728        let child = parent.child();
729        let counter = Arc::new(WakeCounter::default());
730        let waker = Waker::from(Arc::clone(&counter));
731        let mut context = Context::from_waker(&waker);
732        let mut cancelled = Box::pin(child.cancellation().cancelled());
733        assert!(cancelled.as_mut().poll(&mut context).is_pending());
734        assert!(parent.cancellation().cancel_for_resource_exhaustion());
735        assert_eq!(counter.0.load(Ordering::Acquire), 1);
736        assert!(cancelled.as_mut().poll(&mut context).is_ready());
737        assert_eq!(
738            child.checkpoint(),
739            Err(RuntimeCancellationReason::ResourceExhausted)
740        );
741    }
742
743    #[test]
744    fn dropped_cancellation_future_unregisters_from_its_parent_chain() {
745        let parent = RuntimeCancellationToken::new();
746        let child = parent.child();
747        let counter = Arc::new(WakeCounter::default());
748        let waker = Waker::from(Arc::clone(&counter));
749        let mut context = Context::from_waker(&waker);
750        let mut cancelled = Box::pin(child.cancelled());
751
752        assert!(cancelled.as_mut().poll(&mut context).is_pending());
753        drop(cancelled);
754        assert!(parent.cancel());
755        assert_eq!(counter.0.load(Ordering::Acquire), 0);
756    }
757
758    #[test]
759    fn child_preserves_admitted_parallelism() {
760        let context = RuntimeTaskContext::default()
761            .with_admitted_parallelism(NonZeroUsize::new(4).expect("test parallelism is non-zero"))
762            .with_executor_thread_limit(
763                NonZeroUsize::new(2).expect("test thread limit is non-zero"),
764            );
765
766        assert_eq!(context.admitted_parallelism().get(), 4);
767        assert_eq!(context.child().admitted_parallelism().get(), 4);
768        assert_eq!(context.child().executor_thread_limit().unwrap().get(), 2);
769    }
770
771    #[test]
772    fn child_preserves_memory_reservation() {
773        let reservation = RuntimeMemoryReservation::new(128, 32);
774        let context = RuntimeTaskContext::default().with_memory_reservation(reservation);
775
776        assert_eq!(context.memory_reservation(), Some(reservation));
777        assert_eq!(context.child().memory_reservation(), Some(reservation));
778    }
779
780    #[test]
781    fn child_preserves_io_controller_and_permit_releases_slots() {
782        let active_slots = Arc::new(AtomicUsize::new(0));
783        let context = RuntimeTaskContext::default().with_io_wave_controller(Arc::new(
784            RecordingIoController {
785                active_slots: Arc::clone(&active_slots),
786            },
787        ));
788
789        let permit = context
790            .child()
791            .acquire_io_wave(NonZeroUsize::new(2).unwrap())
792            .unwrap()
793            .unwrap();
794        assert_eq!(active_slots.load(Ordering::Acquire), 2);
795        drop(permit);
796        assert_eq!(active_slots.load(Ordering::Acquire), 0);
797    }
798
799    #[test]
800    fn ungoverned_try_acquire_is_immediately_ready() {
801        let acquired = RuntimeTaskContext::default()
802            .try_acquire_io_wave(NonZeroUsize::new(2).unwrap())
803            .unwrap();
804
805        assert!(matches!(acquired, RuntimeIoWaveTryAcquire::Acquired(None)));
806    }
807}
hawdb_core/cancellation.rs

hawdb_core/
cancellation.rs
