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    /// Binds actual allocation ownership to an already admitted task.
406    pub fn with_memory_controller(mut self, controller: Arc<dyn RuntimeMemoryController>) -> Self {
407        self.memory_controller = Some(controller);
408        self
409    }
410
411    /// Reserve before allocating and retain the returned lease until destruction.
412    /// Clones and children share their controller. Ungoverned contexts return
413    /// `None`; a supplied static ceiling still limits an individual allocation.
414    pub fn reserve_working_memory(
415        &self,
416        bytes: u64,
417    ) -> Result<Option<Box<dyn RuntimeMemoryPermit>>, RuntimeMemoryError> {
418        self.checkpoint().map_err(RuntimeMemoryError::Stopped)?;
419        if let Some(reservation) = self.memory_reservation
420            && bytes > reservation.memory_bytes()
421        {
422            return Err(RuntimeMemoryError::ReservationExceeded {
423                requested_bytes: bytes,
424                available_bytes: reservation.memory_bytes(),
425            });
426        }
427        self.memory_controller
428            .as_ref()
429            .map(|controller| {
430                controller.reserve(
431                    bytes,
432                    self.memory_reservation
433                        .map_or(u64::MAX, RuntimeMemoryReservation::memory_bytes),
434                )
435            })
436            .transpose()
437    }
438
439    /// Binds the controller owned by a successful runtime admission.
440    pub fn with_io_wave_controller(mut self, controller: Arc<dyn RuntimeIoWaveController>) -> Self {
441        self.io_wave_controller = Some(controller);
442        self
443    }
444
445    /// Acquires capacity for one storage I/O wave when this is an admitted context.
446    ///
447    /// Raw contexts have no controller and return `Ok(None)`. The returned
448    /// permit must remain live only while physical reads are in flight.
449    pub fn acquire_io_wave(
450        &self,
451        slots: NonZeroUsize,
452    ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
453        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
454        self.io_wave_controller
455            .as_ref()
456            .map(|controller| controller.acquire(slots, self))
457            .transpose()
458    }
459
460    /// Attempts to acquire capacity for one storage I/O wave without blocking.
461    ///
462    /// Raw, ungoverned contexts return `Acquired(None)`. Governed contexts
463    /// return `Pending` while another wave owns the requested capacity.
464    pub fn try_acquire_io_wave(
465        &self,
466        slots: NonZeroUsize,
467    ) -> Result<RuntimeIoWaveTryAcquire, RuntimeIoWaveError> {
468        self.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
469        let Some(controller) = &self.io_wave_controller else {
470            return Ok(RuntimeIoWaveTryAcquire::Acquired(None));
471        };
472        controller.try_acquire(slots, self).map(|permit| {
473            permit.map_or(RuntimeIoWaveTryAcquire::Pending, |permit| {
474                RuntimeIoWaveTryAcquire::Acquired(Some(permit))
475            })
476        })
477    }
478
479    pub fn deadline(&self) -> Option<Instant> {
480        self.deadline
481    }
482
483    pub fn remaining(&self) -> Option<Duration> {
484        self.deadline
485            .map(|deadline| deadline.saturating_duration_since(Instant::now()))
486    }
487
488    pub fn checkpoint(&self) -> Result<(), RuntimeCancellationReason> {
489        if let Some(reason) = self.cancellation.reason() {
490            return Err(reason);
491        }
492        if self
493            .deadline
494            .is_some_and(|deadline| Instant::now() >= deadline)
495        {
496            return Err(RuntimeCancellationReason::DeadlineExceeded);
497        }
498        Ok(())
499    }
500}
501
502impl Default for RuntimeTaskContext {
503    fn default() -> Self {
504        Self::without_deadline(RuntimeCancellationToken::new())
505    }
506}
507
508#[derive(Debug, Clone, Copy, PartialEq, Eq)]
509pub enum RuntimeCancellationReason {
510    Cancelled,
511    DeadlineExceeded,
512    ResourceExhausted,
513}
514
515impl RuntimeCancellationReason {
516    pub const fn as_str(self) -> &'static str {
517        match self {
518            Self::Cancelled => "cancelled",
519            Self::DeadlineExceeded => "deadline_exceeded",
520            Self::ResourceExhausted => "resource_exhausted",
521        }
522    }
523}
524
525impl Display for RuntimeCancellationReason {
526    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
527        formatter.write_str(self.as_str())
528    }
529}
530
531impl Error for RuntimeCancellationReason {}
532
533#[cfg(test)]
534mod tests {
535    use super::*;
536    use std::sync::atomic::AtomicUsize;
537    use std::sync::atomic::Ordering;
538    use std::task::Wake;
539
540    #[derive(Debug)]
541    struct RecordingIoController {
542        active_slots: Arc<AtomicUsize>,
543    }
544
545    #[derive(Debug)]
546    struct RecordingIoPermit {
547        active_slots: Arc<AtomicUsize>,
548        slots: usize,
549    }
550
551    impl Drop for RecordingIoPermit {
552        fn drop(&mut self) {
553            self.active_slots.fetch_sub(self.slots, Ordering::AcqRel);
554        }
555    }
556
557    impl RuntimeIoWaveController for RecordingIoController {
558        fn try_acquire(
559            &self,
560            slots: NonZeroUsize,
561            context: &RuntimeTaskContext,
562        ) -> Result<Option<Box<dyn RuntimeIoWavePermit>>, RuntimeIoWaveError> {
563            self.acquire(slots, context).map(Some)
564        }
565
566        fn acquire(
567            &self,
568            slots: NonZeroUsize,
569            context: &RuntimeTaskContext,
570        ) -> Result<Box<dyn RuntimeIoWavePermit>, RuntimeIoWaveError> {
571            context.checkpoint().map_err(RuntimeIoWaveError::Stopped)?;
572            self.active_slots.fetch_add(slots.get(), Ordering::AcqRel);
573            Ok(Box::new(RecordingIoPermit {
574                active_slots: Arc::clone(&self.active_slots),
575                slots: slots.get(),
576            }))
577        }
578    }
579
580    #[test]
581    fn cancellation_is_shared_across_context_clones() {
582        let token = RuntimeCancellationToken::new();
583        let context = RuntimeTaskContext::without_deadline(token.clone());
584        assert!(context.checkpoint().is_ok());
585        assert!(token.cancel());
586        assert!(!token.cancel());
587        assert_eq!(
588            context.checkpoint(),
589            Err(RuntimeCancellationReason::Cancelled)
590        );
591    }
592
593    #[test]
594    fn expired_deadline_fails_at_a_cooperative_checkpoint() {
595        let context = RuntimeTaskContext::new(
596            RuntimeCancellationToken::new(),
597            Some(Instant::now() - Duration::from_millis(1)),
598        );
599        assert_eq!(
600            context.checkpoint(),
601            Err(RuntimeCancellationReason::DeadlineExceeded)
602        );
603    }
604
605    #[test]
606    fn child_cancellation_is_local_and_parent_cancellation_propagates() {
607        let parent_token = RuntimeCancellationToken::new();
608        let parent = RuntimeTaskContext::without_deadline(parent_token.clone());
609        let child = parent.child();
610
611        assert!(child.cancellation().cancel());
612        assert!(child.checkpoint().is_err());
613        assert!(parent.checkpoint().is_ok());
614
615        let sibling = parent.child();
616        assert!(parent_token.cancel());
617        assert!(parent.checkpoint().is_err());
618        assert!(sibling.checkpoint().is_err());
619    }
620
621    #[test]
622    fn resource_exhaustion_keeps_its_cause_and_is_local_to_the_child() {
623        let parent = RuntimeTaskContext::default();
624        let child = parent.child();
625        let clone = child.clone();
626        let grandchild = child.child();
627        assert!(child.cancellation().cancel_for_resource_exhaustion());
628        assert!(!clone.cancellation().cancel());
629        assert!(!clone.cancellation().cancel_for_resource_exhaustion());
630        for context in [&child, &clone, &grandchild] {
631            assert_eq!(
632                context.checkpoint(),
633                Err(RuntimeCancellationReason::ResourceExhausted)
634            );
635        }
636        assert_eq!(parent.checkpoint(), Ok(()));
637        assert_eq!(parent.child().checkpoint(), Ok(()));
638
639        let cancelled = parent.child();
640        assert!(cancelled.cancellation().cancel());
641        assert!(!cancelled.cancellation().cancel_for_resource_exhaustion());
642        assert_eq!(
643            cancelled.checkpoint(),
644            Err(RuntimeCancellationReason::Cancelled)
645        );
646    }
647
648    #[derive(Default)]
649    struct WakeCounter(AtomicUsize);
650
651    impl Wake for WakeCounter {
652        fn wake(self: Arc<Self>) {
653            self.0.fetch_add(1, Ordering::AcqRel);
654        }
655    }
656
657    #[test]
658    fn cancellation_future_is_woken_by_parent_cancellation() {
659        let parent = RuntimeCancellationToken::new();
660        let child = parent.child();
661        let counter = Arc::new(WakeCounter::default());
662        let waker = Waker::from(Arc::clone(&counter));
663        let mut context = Context::from_waker(&waker);
664        let mut cancelled = Box::pin(child.cancelled());
665
666        assert!(cancelled.as_mut().poll(&mut context).is_pending());
667        assert!(parent.cancel());
668        assert_eq!(counter.0.load(Ordering::Acquire), 1);
669        assert!(cancelled.as_mut().poll(&mut context).is_ready());
670    }
671
672    #[test]
673    fn resource_exhaustion_wakes_waiters_on_descendants() {
674        let parent = RuntimeTaskContext::default();
675        let child = parent.child();
676        let counter = Arc::new(WakeCounter::default());
677        let waker = Waker::from(Arc::clone(&counter));
678        let mut context = Context::from_waker(&waker);
679        let mut cancelled = Box::pin(child.cancellation().cancelled());
680        assert!(cancelled.as_mut().poll(&mut context).is_pending());
681        assert!(parent.cancellation().cancel_for_resource_exhaustion());
682        assert_eq!(counter.0.load(Ordering::Acquire), 1);
683        assert!(cancelled.as_mut().poll(&mut context).is_ready());
684        assert_eq!(
685            child.checkpoint(),
686            Err(RuntimeCancellationReason::ResourceExhausted)
687        );
688    }
689
690    #[test]
691    fn dropped_cancellation_future_unregisters_from_its_parent_chain() {
692        let parent = RuntimeCancellationToken::new();
693        let child = parent.child();
694        let counter = Arc::new(WakeCounter::default());
695        let waker = Waker::from(Arc::clone(&counter));
696        let mut context = Context::from_waker(&waker);
697        let mut cancelled = Box::pin(child.cancelled());
698
699        assert!(cancelled.as_mut().poll(&mut context).is_pending());
700        drop(cancelled);
701        assert!(parent.cancel());
702        assert_eq!(counter.0.load(Ordering::Acquire), 0);
703    }
704
705    #[test]
706    fn child_preserves_admitted_parallelism() {
707        let context = RuntimeTaskContext::default()
708            .with_admitted_parallelism(NonZeroUsize::new(4).expect("test parallelism is non-zero"))
709            .with_executor_thread_limit(
710                NonZeroUsize::new(2).expect("test thread limit is non-zero"),
711            );
712
713        assert_eq!(context.admitted_parallelism().get(), 4);
714        assert_eq!(context.child().admitted_parallelism().get(), 4);
715        assert_eq!(context.child().executor_thread_limit().unwrap().get(), 2);
716    }
717
718    #[test]
719    fn child_preserves_memory_reservation() {
720        let reservation = RuntimeMemoryReservation::new(128, 32);
721        let context = RuntimeTaskContext::default().with_memory_reservation(reservation);
722
723        assert_eq!(context.memory_reservation(), Some(reservation));
724        assert_eq!(context.child().memory_reservation(), Some(reservation));
725    }
726
727    #[test]
728    fn child_preserves_io_controller_and_permit_releases_slots() {
729        let active_slots = Arc::new(AtomicUsize::new(0));
730        let context = RuntimeTaskContext::default().with_io_wave_controller(Arc::new(
731            RecordingIoController {
732                active_slots: Arc::clone(&active_slots),
733            },
734        ));
735
736        let permit = context
737            .child()
738            .acquire_io_wave(NonZeroUsize::new(2).unwrap())
739            .unwrap()
740            .unwrap();
741        assert_eq!(active_slots.load(Ordering::Acquire), 2);
742        drop(permit);
743        assert_eq!(active_slots.load(Ordering::Acquire), 0);
744    }
745
746    #[test]
747    fn ungoverned_try_acquire_is_immediately_ready() {
748        let acquired = RuntimeTaskContext::default()
749            .try_acquire_io_wave(NonZeroUsize::new(2).unwrap())
750            .unwrap();
751
752        assert!(matches!(acquired, RuntimeIoWaveTryAcquire::Acquired(None)));
753    }
754}
hawdb_core/cancellation.rs

hawdb_core/
cancellation.rs
