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;
270pub use memory::{RuntimeMemoryController, RuntimeMemoryError, RuntimeMemoryPermit};
271
272#[derive(Debug, Clone)]
273pub struct RuntimeTaskContext {
274    cancellation: RuntimeCancellationToken,
275    deadline: Option<Instant>,
276    admitted_parallelism: NonZeroUsize,
277    executor_thread_limit: Option<NonZeroUsize>,
278    memory_reservation: Option<RuntimeMemoryReservation>,
279    memory_controller: Option<Arc<dyn RuntimeMemoryController>>,
280    io_wave_controller: Option<Arc<dyn RuntimeIoWaveController>>,
281}
282
283/// Memory already reserved for one task by the runtime governor.
284#[derive(Debug, Clone, Copy, PartialEq, Eq)]
285pub struct RuntimeMemoryReservation {
286    memory_bytes: u64,
287    result_bytes: u64,
288}
289
290impl RuntimeMemoryReservation {
291    pub const fn new(memory_bytes: u64, result_bytes: u64) -> Self {
292        Self {
293            memory_bytes,
294            result_bytes,
295        }
296    }
297
298    pub const fn memory_bytes(self) -> u64 {
299        self.memory_bytes
300    }
301
302    pub const fn result_bytes(self) -> u64 {
303        self.result_bytes
304    }
305
306    pub const fn intersect(self, other: Self) -> Self {
307        Self {
308            memory_bytes: if self.memory_bytes < other.memory_bytes {
309                self.memory_bytes
310            } else {
311                other.memory_bytes
312            },
313            result_bytes: if self.result_bytes < other.result_bytes {
314                self.result_bytes
315            } else {
316                other.result_bytes
317            },
318        }
319    }
320}
321
322impl RuntimeTaskContext {
323    pub fn new(cancellation: RuntimeCancellationToken, deadline: Option<Instant>) -> Self {
324        Self {
325            cancellation,
326            deadline,
327            admitted_parallelism: NonZeroUsize::MIN,
328            executor_thread_limit: None,
329            memory_reservation: None,
330            memory_controller: None,
331            io_wave_controller: None,
332        }
333    }
334
335    pub fn without_deadline(cancellation: RuntimeCancellationToken) -> Self {
336        Self::new(cancellation, None)
337    }
338
339    pub fn with_timeout(timeout: Duration) -> Self {
340        Self::new(
341            RuntimeCancellationToken::new(),
342            Instant::now().checked_add(timeout),
343        )
344    }
345
346    pub fn cancellation(&self) -> &RuntimeCancellationToken {
347        &self.cancellation
348    }
349
350    pub fn child(&self) -> Self {
351        Self {
352            cancellation: self.cancellation.child(),
353            deadline: self.deadline,
354            admitted_parallelism: self.admitted_parallelism,
355            executor_thread_limit: self.executor_thread_limit,
356            memory_reservation: self.memory_reservation,
357            memory_controller: self.memory_controller.clone(),
358            io_wave_controller: self.io_wave_controller.clone(),
359        }
360    }
361
362    /// Carries the CPU parallelism already reserved by the runtime governor.
363    ///
364    /// This is an execution ceiling, not a request to create worker threads.
365    /// Executors must still apply their operator memory and input-size limits.
366    pub fn with_admitted_parallelism(mut self, admitted_parallelism: NonZeroUsize) -> Self {
367        self.admitted_parallelism = admitted_parallelism;
368        self
369    }
370
371    pub fn admitted_parallelism(&self) -> NonZeroUsize {
372        self.admitted_parallelism
373    }
374
375    /// Carries the physical executor-thread ceiling derived by the runtime governor.
376    ///
377    /// This is distinct from [`Self::admitted_parallelism`], which is the CPU
378    /// reservation for one task. An absent limit denotes an ungoverned library call.
379    pub fn with_executor_thread_limit(mut self, limit: NonZeroUsize) -> Self {
380        self.executor_thread_limit = Some(
381            self.executor_thread_limit
382                .map_or(limit, |current| current.min(limit)),
383        );
384        self
385    }
386
387    pub fn executor_thread_limit(&self) -> Option<NonZeroUsize> {
388        self.executor_thread_limit
389    }
390
391    /// Carries the memory already reserved by the runtime governor.
392    ///
393    /// Governed executors must use this reservation instead of a static
394    /// per-query configuration limit. An absent reservation denotes an
395    /// ungoverned library call and preserves the configured fallback.
396    pub fn with_memory_reservation(mut self, reservation: RuntimeMemoryReservation) -> Self {
397        self.memory_reservation = Some(reservation);
398        self
399    }
400
401    pub fn memory_reservation(&self) -> Option<RuntimeMemoryReservation> {
402        self.memory_reservation
403    }
404
405    /// Whether execution carries dynamic resource controllers rather than only
406    /// a static ceiling. Synchronous adapters must preserve these controllers
407    /// and must not replace an originating governor admission with a local cap.
408    pub fn has_resource_controllers(&self) -> bool {
409        self.memory_controller.is_some() || self.io_wave_controller.is_some()
410    }
411
412    /// Binds actual allocation ownership to an already admitted task.
413    pub fn with_memory_controller(mut self, controller: Arc<dyn RuntimeMemoryController>) -> Self {
414        self.memory_controller = Some(controller);
415        self
416    }
417
418    /// Reserve before allocating and retain the returned lease until destruction.
419    /// Clones and children share their controller. Ungoverned contexts return
420    /// `None`; a supplied static ceiling still limits an individual allocation.
421    pub fn reserve_working_memory(
422        &self,
423        bytes: u64,
424    ) -> Result<Option<Box<dyn RuntimeMemoryPermit>>, RuntimeMemoryError> {
425        self.checkpoint().map_err(RuntimeMemoryError::Stopped)?;
426        if let Some(reservation) = self.memory_reservation
427            && bytes > reservation.memory_bytes()
428        {
429            return Err(RuntimeMemoryError::ReservationExceeded {
430                requested_bytes: bytes,
431                available_bytes: reservation.memory_bytes(),
432            });
433        }
434        self.memory_controller
435            .as_ref()
436            .map(|controller| {
437                controller.reserve(
438                    bytes,
439                    self.memory_reservation
440                        .map_or(u64::MAX, RuntimeMemoryReservation::memory_bytes),
441                )
442            })
443            .transpose()
444    }
445
446    /// Binds the controller owned by a successful runtime admission.
447    pub fn with_io_wave_controller(mut self, controller: Arc<dyn RuntimeIoWaveController>) -> Self {
448        self.io_wave_controller = Some(controller);
449        self
450    }
451
452    /// Acquires capacity for one storage I/O wave when this is an admitted context.
453    ///
454    /// Raw contexts have no controller and return `Ok(None)`. The returned
455    /// permit must remain live only while physical reads are in flight.
456    pub fn acquire_io_wave(
457        &self,
458        slots: NonZeroUsize,
459    ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
460        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
461        self.io_wave_controller
462            .as_ref()
463            .map(|controller| controller.acquire(slots, self))
464            .transpose()
465    }
466
467    /// Attempts to acquire capacity for one storage I/O wave without blocking.
468    ///
469    /// Raw, ungoverned contexts return `Acquired(None)`. Governed contexts
470    /// return `Pending` while another wave owns the requested capacity.
471    pub fn try_acquire_io_wave(
472        &self,
473        slots: NonZeroUsize,
474    ) -> Result<RuntimeIoWaveTryAcquire, RuntimeIoWaveError> {
475        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
476        let Some(controller) = &self.io_wave_controller else {
477            return Ok(RuntimeIoWaveTryAcquire::Acquired(None));
478        };
479        controller.try_acquire(slots, self).map(|permit| {
480            permit.map_or(RuntimeIoWaveTryAcquire::Pending, |permit| {
481                RuntimeIoWaveTryAcquire::Acquired(Some(permit))
482            })
483        })
484    }
485
486    pub fn deadline(&self) -> Option<Instant> {
487        self.deadline
488    }
489
490    pub fn remaining(&self) -> Option<Duration> {
491        self.deadline
492            .map(|deadline| deadline.saturating_duration_since(Instant::now()))
493    }
494
495    pub fn checkpoint(&self) -> Result<(), RuntimeCancellationReason> {
496        if let Some(reason) = self.cancellation.reason() {
497            return Err(reason);
498        }
499        if self
500            .deadline
501            .is_some_and(|deadline| Instant::now() >= deadline)
502        {
503            return Err(RuntimeCancellationReason::DeadlineExceeded);
504        }
505        Ok(())
506    }
507}
508
509impl Default for RuntimeTaskContext {
510    fn default() -> Self {
511        Self::without_deadline(RuntimeCancellationToken::new())
512    }
513}
514
515#[derive(Debug, Clone, Copy, PartialEq, Eq)]
516pub enum RuntimeCancellationReason {
517    Cancelled,
518    DeadlineExceeded,
519    ResourceExhausted,
520}
521
522impl RuntimeCancellationReason {
523    pub const fn as_str(self) -> &'static str {
524        match self {
525            Self::Cancelled => "cancelled",
526            Self::DeadlineExceeded => "deadline_exceeded",
527            Self::ResourceExhausted => "resource_exhausted",
528        }
529    }
530}
531
532impl Display for RuntimeCancellationReason {
533    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
534        formatter.write_str(self.as_str())
535    }
536}
537
538impl Error for RuntimeCancellationReason {}
539
540#[cfg(test)]
541mod tests {
542    use super::*;
543    use std::sync::atomic::AtomicUsize;
544    use std::sync::atomic::Ordering;
545    use std::task::Wake;
546
547    #[derive(Debug)]
548    struct RecordingIoController {
549        active_slots: Arc<AtomicUsize>,
550    }
551
552    #[derive(Debug)]
553    struct RecordingIoPermit {
554        active_slots: Arc<AtomicUsize>,
555        slots: usize,
556    }
557
558    impl Drop for RecordingIoPermit {
559        fn drop(&mut self) {
560            self.active_slots.fetch_sub(self.slots, Ordering::AcqRel);
561        }
562    }
563
564    impl RuntimeIoWaveController for RecordingIoController {
565        fn try_acquire(
566            &self,
567            slots: NonZeroUsize,
568            context: &RuntimeTaskContext,
569        ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
570            self.acquire(slots, context).map(Some)
571        }
572
573        fn acquire(
574            &self,
575            slots: NonZeroUsize,
576            context: &RuntimeTaskContext,
577        ) -> Result<Box<dyn RuntimeIoWavePermit>, RuntimeIoWaveError> {
578            context.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
579            self.active_slots.fetch_add(slots.get(), Ordering::AcqRel);
580            Ok(Box::new(RecordingIoPermit {
581                active_slots: Arc::clone(&self.active_slots),
582                slots: slots.get(),
583            }))
584        }
585    }
586
587    #[test]
588    fn cancellation_is_shared_across_context_clones() {
589        let token = RuntimeCancellationToken::new();
590        let context = RuntimeTaskContext::without_deadline(token.clone());
591        assert!(context.checkpoint().is_ok());
592        assert!(token.cancel());
593        assert!(!token.cancel());
594        assert_eq!(
595            context.checkpoint(),
596            Err(RuntimeCancellationReason::Cancelled)
597        );
598    }
599
600    #[test]
601    fn expired_deadline_fails_at_a_cooperative_checkpoint() {
602        let context = RuntimeTaskContext::new(
603            RuntimeCancellationToken::new(),
604            Some(Instant::now() - Duration::from_millis(1)),
605        );
606        assert_eq!(
607            context.checkpoint(),
608            Err(RuntimeCancellationReason::DeadlineExceeded)
609        );
610    }
611
612    #[test]
613    fn child_cancellation_is_local_and_parent_cancellation_propagates() {
614        let parent_token = RuntimeCancellationToken::new();
615        let parent = RuntimeTaskContext::without_deadline(parent_token.clone());
616        let child = parent.child();
617
618        assert!(child.cancellation().cancel());
619        assert!(child.checkpoint().is_err());
620        assert!(parent.checkpoint().is_ok());
621
622        let sibling = parent.child();
623        assert!(parent_token.cancel());
624        assert!(parent.checkpoint().is_err());
625        assert!(sibling.checkpoint().is_err());
626    }
627
628    #[test]
629    fn resource_exhaustion_keeps_its_cause_and_is_local_to_the_child() {
630        let parent = RuntimeTaskContext::default();
631        let child = parent.child();
632        let clone = child.clone();
633        let grandchild = child.child();
634        assert!(child.cancellation().cancel_for_resource_exhaustion());
635        assert!(!clone.cancellation().cancel());
636        assert!(!clone.cancellation().cancel_for_resource_exhaustion());
637        for context in [&child, &clone, &grandchild] {
638            assert_eq!(
639                context.checkpoint(),
640                Err(RuntimeCancellationReason::ResourceExhausted)
641            );
642        }
643        assert_eq!(parent.checkpoint(), Ok(()));
644        assert_eq!(parent.child().checkpoint(), Ok(()));
645
646        let cancelled = parent.child();
647        assert!(cancelled.cancellation().cancel());
648        assert!(!cancelled.cancellation().cancel_for_resource_exhaustion());
649        assert_eq!(
650            cancelled.checkpoint(),
651            Err(RuntimeCancellationReason::Cancelled)
652        );
653    }
654
655    #[derive(Default)]
656    struct WakeCounter(AtomicUsize);
657
658    impl Wake for WakeCounter {
659        fn wake(self: Arc<Self>) {
660            self.0.fetch_add(1, Ordering::AcqRel);
661        }
662    }
663
664    #[test]
665    fn cancellation_future_is_woken_by_parent_cancellation() {
666        let parent = RuntimeCancellationToken::new();
667        let child = parent.child();
668        let counter = Arc::new(WakeCounter::default());
669        let waker = Waker::from(Arc::clone(&counter));
670        let mut context = Context::from_waker(&waker);
671        let mut cancelled = Box::pin(child.cancelled());
672
673        assert!(cancelled.as_mut().poll(&mut context).is_pending());
674        assert!(parent.cancel());
675        assert_eq!(counter.0.load(Ordering::Acquire), 1);
676        assert!(cancelled.as_mut().poll(&mut context).is_ready());
677    }
678
679    #[test]
680    fn resource_exhaustion_wakes_waiters_on_descendants() {
681        let parent = RuntimeTaskContext::default();
682        let child = parent.child();
683        let counter = Arc::new(WakeCounter::default());
684        let waker = Waker::from(Arc::clone(&counter));
685        let mut context = Context::from_waker(&waker);
686        let mut cancelled = Box::pin(child.cancellation().cancelled());
687        assert!(cancelled.as_mut().poll(&mut context).is_pending());
688        assert!(parent.cancellation().cancel_for_resource_exhaustion());
689        assert_eq!(counter.0.load(Ordering::Acquire), 1);
690        assert!(cancelled.as_mut().poll(&mut context).is_ready());
691        assert_eq!(
692            child.checkpoint(),
693            Err(RuntimeCancellationReason::ResourceExhausted)
694        );
695    }
696
697    #[test]
698    fn dropped_cancellation_future_unregisters_from_its_parent_chain() {
699        let parent = RuntimeCancellationToken::new();
700        let child = parent.child();
701        let counter = Arc::new(WakeCounter::default());
702        let waker = Waker::from(Arc::clone(&counter));
703        let mut context = Context::from_waker(&waker);
704        let mut cancelled = Box::pin(child.cancelled());
705
706        assert!(cancelled.as_mut().poll(&mut context).is_pending());
707        drop(cancelled);
708        assert!(parent.cancel());
709        assert_eq!(counter.0.load(Ordering::Acquire), 0);
710    }
711
712    #[test]
713    fn child_preserves_admitted_parallelism() {
714        let context = RuntimeTaskContext::default()
715            .with_admitted_parallelism(NonZeroUsize::new(4).expect("test parallelism is non-zero"))
716            .with_executor_thread_limit(
717                NonZeroUsize::new(2).expect("test thread limit is non-zero"),
718            );
719
720        assert_eq!(context.admitted_parallelism().get(), 4);
721        assert_eq!(context.child().admitted_parallelism().get(), 4);
722        assert_eq!(context.child().executor_thread_limit().unwrap().get(), 2);
723    }
724
725    #[test]
726    fn child_preserves_memory_reservation() {
727        let reservation = RuntimeMemoryReservation::new(128, 32);
728        let context = RuntimeTaskContext::default().with_memory_reservation(reservation);
729
730        assert_eq!(context.memory_reservation(), Some(reservation));
731        assert_eq!(context.child().memory_reservation(), Some(reservation));
732    }
733
734    #[test]
735    fn child_preserves_io_controller_and_permit_releases_slots() {
736        let active_slots = Arc::new(AtomicUsize::new(0));
737        let context = RuntimeTaskContext::default().with_io_wave_controller(Arc::new(
738            RecordingIoController {
739                active_slots: Arc::clone(&active_slots),
740            },
741        ));
742
743        let permit = context
744            .child()
745            .acquire_io_wave(NonZeroUsize::new(2).unwrap())
746            .unwrap()
747            .unwrap();
748        assert_eq!(active_slots.load(Ordering::Acquire), 2);
749        drop(permit);
750        assert_eq!(active_slots.load(Ordering::Acquire), 0);
751    }
752
753    #[test]
754    fn ungoverned_try_acquire_is_immediately_ready() {
755        let acquired = RuntimeTaskContext::default()
756            .try_acquire_io_wave(NonZeroUsize::new(2).unwrap())
757            .unwrap();
758
759        assert!(matches!(acquired, RuntimeIoWaveTryAcquire::Acquired(None)));
760    }
761}
hawdb_core/cancellation.rs

hawdb_core/
cancellation.rs
