cadmpeg_core/decode/
budget.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Unified resource accounting for one decode session.
3
4use std::cell::Cell;
5use std::num::NonZeroU64;
6
7use crate::CodecError;
8
9use super::error::{ResourceDimension, ResourceFailure, ResourceLimit};
10use super::policy::{
11    DecodePolicy, DECOMPRESSED_TOTAL_BASE, DECOMPRESSED_TOTAL_PER_INPUT_BYTE, MATERIALIZED_BASE,
12    MATERIALIZED_PER_INPUT_BYTE, RETAINED_BASE, RETAINED_PER_INPUT_BYTE,
13};
14use super::view::u64_from_index;
15
16#[derive(Debug)]
17pub(super) struct DecodeBudget {
18    /// Policy applied to this budget.
19    policy: DecodePolicy,
20    input_bytes: Cell<u64>,
21    decompressed: Cell<u64>,
22    materialized: Cell<u64>,
23    retained: Cell<u64>,
24    scoped_storage: Cell<Option<u64>>,
25    entities: Cell<u64>,
26    collection_items: Cell<u64>,
27    recursion_depth: Cell<u64>,
28    work: Cell<u64>,
29    fuse: Cell<Option<ResourceLimit>>,
30}
31
32impl DecodeBudget {
33    /// Returns the policy applied to this budget.
34    pub(super) fn policy(&self) -> &DecodePolicy {
35        &self.policy
36    }
37
38    pub(super) fn new(policy: DecodePolicy, input_bytes: u64) -> Self {
39        Self {
40            policy,
41            input_bytes: Cell::new(input_bytes),
42            decompressed: Cell::new(0),
43            materialized: Cell::new(0),
44            retained: Cell::new(0),
45            scoped_storage: Cell::new(None),
46            entities: Cell::new(0),
47            collection_items: Cell::new(0),
48            recursion_depth: Cell::new(0),
49            work: Cell::new(0),
50            fuse: Cell::new(None),
51        }
52    }
53
54    pub(super) fn fused(&self) -> Option<ResourceLimit> {
55        self.fuse.get()
56    }
57
58    pub(super) fn decompression_allowance(&self) -> u64 {
59        let policy_limit = self.policy.limits.max_decompressed_bytes_total;
60        let Some(proportional) = DECOMPRESSED_TOTAL_PER_INPUT_BYTE
61            .checked_mul(self.input_bytes.get())
62            .and_then(|bytes| DECOMPRESSED_TOTAL_BASE.checked_add(bytes))
63        else {
64            return policy_limit;
65        };
66        policy_limit.min(proportional)
67    }
68
69    pub(super) fn input_bytes(&self) -> u64 {
70        self.input_bytes.get()
71    }
72
73    pub(super) fn charge_input(
74        &self,
75        amount: u64,
76        operation: &'static str,
77    ) -> Result<(), CodecError> {
78        self.charge(
79            ResourceDimension::InputBytes,
80            &self.input_bytes,
81            self.policy.limits.max_input_bytes,
82            amount,
83            operation,
84        )
85        .map_err(Into::into)
86    }
87
88    pub(super) fn charge_decompressed(
89        &self,
90        amount: u64,
91        operation: &'static str,
92    ) -> Result<(), CodecError> {
93        self.charge(
94            ResourceDimension::DecompressedBytes,
95            &self.decompressed,
96            self.decompression_allowance(),
97            amount,
98            operation,
99        )
100        .map_err(Into::into)
101    }
102
103    pub(super) fn decompressed_used(&self) -> u64 {
104        self.decompressed.get()
105    }
106
107    fn materialized_allowance(&self) -> u64 {
108        let policy_limit = self.policy.limits.max_materialized_bytes;
109        let Some(proportional) = MATERIALIZED_PER_INPUT_BYTE
110            .checked_mul(self.input_bytes.get())
111            .and_then(|bytes| MATERIALIZED_BASE.checked_add(bytes))
112        else {
113            return policy_limit;
114        };
115        policy_limit.min(proportional)
116    }
117
118    fn retained_allowance(&self) -> u64 {
119        let policy_limit = self.policy.limits.max_retained_bytes;
120        let Some(proportional) = RETAINED_PER_INPUT_BYTE
121            .checked_mul(self.input_bytes.get())
122            .and_then(|bytes| RETAINED_BASE.checked_add(bytes))
123        else {
124            return policy_limit;
125        };
126        policy_limit.min(proportional)
127    }
128
129    fn charge(
130        &self,
131        dimension: ResourceDimension,
132        used: &Cell<u64>,
133        limit: u64,
134        amount: u64,
135        operation: &'static str,
136    ) -> Result<(), ResourceLimit> {
137        if let Some(resource) = self.fuse.get() {
138            return Err(resource);
139        }
140        let before = used.get();
141        #[cfg(feature = "test-support")]
142        let limit = super::test_support::charge_limit(
143            &self.policy.limits,
144            dimension,
145            operation,
146            before,
147            amount,
148        )
149        .map_or(limit, |probe| limit.min(probe));
150        if limit
151            .checked_sub(before)
152            .is_none_or(|remaining| amount > remaining)
153        {
154            return Err(self.refuse_limit(
155                dimension,
156                ResourceFailure::BudgetExceeded,
157                limit,
158                before,
159                amount,
160                operation,
161            ));
162        }
163        used.set(before + amount);
164        Ok(())
165    }
166
167    pub(super) fn refuse_limit(
168        &self,
169        dimension: ResourceDimension,
170        reason: ResourceFailure,
171        limit: u64,
172        used: u64,
173        additional: u64,
174        operation: &'static str,
175    ) -> ResourceLimit {
176        let resource = ResourceLimit {
177            dimension,
178            reason,
179            limit,
180            used,
181            additional,
182            operation,
183        };
184        self.fuse.set(Some(resource));
185        resource
186    }
187
188    pub(super) fn refuse(
189        &self,
190        dimension: ResourceDimension,
191        reason: ResourceFailure,
192        limit: u64,
193        used: u64,
194        additional: u64,
195        operation: &'static str,
196    ) -> CodecError {
197        self.refuse_limit(dimension, reason, limit, used, additional, operation)
198            .into()
199    }
200
201    pub(super) fn reserve_scoped(
202        &self,
203        bytes: u64,
204        operation: &'static str,
205    ) -> Result<ScopedReservation<'_>, CodecError> {
206        self.reserve_scoped_limit(bytes, operation)
207            .map_err(Into::into)
208    }
209
210    pub(super) fn reserve_scoped_limit(
211        &self,
212        bytes: u64,
213        operation: &'static str,
214    ) -> Result<ScopedReservation<'_>, ResourceLimit> {
215        self.charge(
216            ResourceDimension::MaterializedBytes,
217            &self.materialized,
218            self.materialized_allowance(),
219            bytes,
220            operation,
221        )?;
222        Ok(ScopedReservation {
223            budget: self,
224            bytes,
225            operation,
226        })
227    }
228
229    /// Report allocator refusal after a scoped-byte reservation was recorded.
230    pub(super) fn scoped_allocation_failed(
231        &self,
232        charged: u64,
233        operation: &'static str,
234    ) -> CodecError {
235        self.scoped_allocation_failed_limit(charged, operation)
236            .into()
237    }
238
239    pub(super) fn scoped_allocation_failed_limit(
240        &self,
241        charged: u64,
242        operation: &'static str,
243    ) -> ResourceLimit {
244        self.refuse_limit(
245            ResourceDimension::MaterializedBytes,
246            ResourceFailure::AllocationFailed,
247            self.materialized_allowance(),
248            self.materialized.get() - charged,
249            charged,
250            operation,
251        )
252    }
253
254    pub(super) fn scoped_size_overflow_limit(&self, operation: &'static str) -> ResourceLimit {
255        self.refuse_limit(
256            ResourceDimension::MaterializedBytes,
257            ResourceFailure::BudgetExceeded,
258            self.materialized_allowance(),
259            self.materialized.get(),
260            u64::MAX,
261            operation,
262        )
263    }
264
265    pub(super) fn charge_retained(
266        &self,
267        bytes: u64,
268        operation: &'static str,
269    ) -> Result<(), CodecError> {
270        self.charge_retained_limit(bytes, operation)
271            .map_err(Into::into)
272    }
273
274    pub(super) fn charge_retained_limit(
275        &self,
276        bytes: u64,
277        operation: &'static str,
278    ) -> Result<(), ResourceLimit> {
279        if let Some(current) = self.scoped_storage.get() {
280            let total = current.checked_add(bytes).ok_or_else(|| {
281                self.refuse_limit(
282                    ResourceDimension::MaterializedBytes,
283                    ResourceFailure::BudgetExceeded,
284                    self.materialized_allowance(),
285                    current,
286                    bytes,
287                    operation,
288                )
289            })?;
290            self.charge(
291                ResourceDimension::MaterializedBytes,
292                &self.materialized,
293                self.materialized_allowance(),
294                bytes,
295                operation,
296            )?;
297            self.scoped_storage.set(Some(total));
298            return Ok(());
299        }
300        self.charge(
301            ResourceDimension::RetainedBytes,
302            &self.retained,
303            self.retained_allowance(),
304            bytes,
305            operation,
306        )
307    }
308
309    pub(super) fn retained_size_overflow_limit(&self, operation: &'static str) -> ResourceLimit {
310        let (dimension, limit, used) = if self.scoped_storage.get().is_some() {
311            (
312                ResourceDimension::MaterializedBytes,
313                self.materialized_allowance(),
314                self.materialized.get(),
315            )
316        } else {
317            (
318                ResourceDimension::RetainedBytes,
319                self.retained_allowance(),
320                self.retained.get(),
321            )
322        };
323        self.refuse_limit(
324            dimension,
325            ResourceFailure::BudgetExceeded,
326            limit,
327            used,
328            u64::MAX,
329            operation,
330        )
331    }
332
333    /// Report allocator refusal after a retained charge was already recorded.
334    pub(super) fn retained_allocation_failed(
335        &self,
336        charged: u64,
337        operation: &'static str,
338    ) -> CodecError {
339        self.retained_allocation_failed_limit(charged, operation)
340            .into()
341    }
342
343    pub(super) fn retained_allocation_failed_limit(
344        &self,
345        charged: u64,
346        operation: &'static str,
347    ) -> ResourceLimit {
348        let (dimension, current, limit) = if self.scoped_storage.get().is_some() {
349            (
350                ResourceDimension::MaterializedBytes,
351                self.materialized.get(),
352                self.materialized_allowance(),
353            )
354        } else {
355            (
356                ResourceDimension::RetainedBytes,
357                self.retained.get(),
358                self.retained_allowance(),
359            )
360        };
361        let prior = match current.checked_sub(charged) {
362            Some(prior) => prior,
363            None => current,
364        };
365        self.refuse_limit(
366            dimension,
367            ResourceFailure::AllocationFailed,
368            limit,
369            prior,
370            charged,
371            operation,
372        )
373    }
374
375    pub(super) fn charge_entities(
376        &self,
377        count: u64,
378        operation: &'static str,
379    ) -> Result<(), CodecError> {
380        self.charge(
381            ResourceDimension::Entities,
382            &self.entities,
383            self.policy.limits.max_entities,
384            count,
385            operation,
386        )
387        .map_err(Into::into)
388    }
389
390    pub(super) fn charge_collection_items(
391        &self,
392        count: u64,
393        operation: &'static str,
394    ) -> Result<(), CodecError> {
395        self.charge_collection_items_limit(count, operation)
396            .map_err(Into::into)
397    }
398
399    pub(super) fn charge_collection_items_limit(
400        &self,
401        count: u64,
402        operation: &'static str,
403    ) -> Result<(), ResourceLimit> {
404        self.charge(
405            ResourceDimension::CollectionItems,
406            &self.collection_items,
407            self.policy.limits.max_collection_items,
408            count,
409            operation,
410        )
411    }
412
413    pub(super) fn collection_allocation_failed(
414        &self,
415        charged: u64,
416        operation: &'static str,
417    ) -> CodecError {
418        self.collection_allocation_failed_limit(charged, operation)
419            .into()
420    }
421
422    pub(super) fn collection_allocation_failed_limit(
423        &self,
424        charged: u64,
425        operation: &'static str,
426    ) -> ResourceLimit {
427        let current = self.collection_items.get();
428        let prior = match current.checked_sub(charged) {
429            Some(prior) => prior,
430            None => current,
431        };
432        self.refuse_limit(
433            ResourceDimension::CollectionItems,
434            ResourceFailure::AllocationFailed,
435            self.policy.limits.max_collection_items,
436            prior,
437            charged,
438            operation,
439        )
440    }
441
442    pub(super) fn storage_scope(&self, operation: &'static str) -> StorageScope<'_> {
443        let prior = self.scoped_storage.replace(Some(0));
444        StorageScope {
445            budget: self,
446            prior,
447            operation,
448        }
449    }
450
451    pub(super) fn enter_nested(
452        &self,
453        operation: &'static str,
454    ) -> Result<DepthGuard<'_>, ResourceLimit> {
455        self.charge(
456            ResourceDimension::RecursionDepth,
457            &self.recursion_depth,
458            self.policy.limits.max_recursion_depth,
459            1,
460            operation,
461        )?;
462        Ok(DepthGuard { budget: self })
463    }
464
465    pub(super) fn charge_work(
466        &self,
467        units: u64,
468        operation: &'static str,
469    ) -> Result<(), CodecError> {
470        self.charge_work_limit(units, operation).map_err(Into::into)
471    }
472
473    pub(super) fn charge_work_limit(
474        &self,
475        units: u64,
476        operation: &'static str,
477    ) -> Result<(), ResourceLimit> {
478        self.charge(
479            ResourceDimension::WorkUnits,
480            &self.work,
481            self.policy.limits.max_work_units,
482            units,
483            operation,
484        )
485    }
486}
487
488pub(super) struct StorageScope<'a> {
489    budget: &'a DecodeBudget,
490    prior: Option<u64>,
491    operation: &'static str,
492}
493
494impl<'a> StorageScope<'a> {
495    pub(super) fn finish(self) -> ScopedReservation<'a> {
496        let bytes = self
497            .budget
498            .scoped_storage
499            .replace(Some(0))
500            .map_or(0, std::convert::identity);
501        ScopedReservation {
502            budget: self.budget,
503            bytes,
504            operation: self.operation,
505        }
506    }
507}
508
509impl Drop for StorageScope<'_> {
510    fn drop(&mut self) {
511        let bytes = self
512            .budget
513            .scoped_storage
514            .replace(self.prior)
515            .map_or(0, std::convert::identity);
516        drop(ScopedReservation {
517            budget: self.budget,
518            bytes,
519            operation: self.operation,
520        });
521    }
522}
523
524/// A temporary materialization charged until it is dropped or committed.
525#[derive(Debug)]
526pub struct ScopedReservation<'a> {
527    budget: &'a DecodeBudget,
528    bytes: u64,
529    operation: &'static str,
530}
531
532impl ScopedReservation<'_> {
533    /// Account copied storage in this live temporary reservation.
534    pub fn with_storage<T, E: From<CodecError>>(
535        &mut self,
536        build: impl FnOnce() -> Result<T, E>,
537    ) -> Result<T, E> {
538        self.with_storage_error(build, |limit| E::from(CodecError::ResourceLimit(limit)))
539    }
540
541    /// Account temporary storage with a resource-only refusal channel.
542    pub fn with_storage_limit<T>(
543        &mut self,
544        build: impl FnOnce() -> Result<T, ResourceLimit>,
545    ) -> Result<T, ResourceLimit> {
546        self.with_storage_error(build, |limit| limit)
547    }
548
549    fn with_storage_error<T, E>(
550        &mut self,
551        build: impl FnOnce() -> Result<T, E>,
552        failure: impl FnOnce(ResourceLimit) -> E,
553    ) -> Result<T, E> {
554        let scope = self.budget.storage_scope(self.operation);
555        let value = build();
556        let mut storage = scope.finish();
557        let Some(bytes) = self.bytes.checked_add(storage.bytes) else {
558            return value.and_then(|_| {
559                Err(failure(self.budget.refuse_limit(
560                    ResourceDimension::MaterializedBytes,
561                    ResourceFailure::BudgetExceeded,
562                    self.budget.materialized_allowance(),
563                    self.bytes,
564                    storage.bytes,
565                    self.operation,
566                )))
567            });
568        };
569        self.bytes = bytes;
570        storage.bytes = 0;
571        value
572    }
573
574    /// Increases live temporary storage and returns the resource refusal.
575    pub fn grow_limit(&mut self, bytes: u64) -> Result<(), ResourceLimit> {
576        self.budget.charge(
577            ResourceDimension::MaterializedBytes,
578            &self.budget.materialized,
579            self.budget.materialized_allowance(),
580            bytes,
581            self.operation,
582        )?;
583        self.bytes = self.bytes.checked_add(bytes).ok_or_else(|| {
584            self.budget.refuse_limit(
585                ResourceDimension::MaterializedBytes,
586                ResourceFailure::BudgetExceeded,
587                self.budget.materialized_allowance(),
588                self.bytes,
589                bytes,
590                self.operation,
591            )
592        })?;
593        Ok(())
594    }
595
596    /// Increases the live temporary reservation.
597    pub fn grow(&mut self, bytes: u64) -> Result<(), CodecError> {
598        self.grow_limit(bytes).map_err(Into::into)
599    }
600
601    /// Converts the temporary reservation into session-retained bytes.
602    pub fn commit(self) -> Result<(), CodecError> {
603        self.budget.charge_retained(self.bytes, self.operation)
604    }
605}
606
607impl Drop for ScopedReservation<'_> {
608    fn drop(&mut self) {
609        let Some(remaining) = self.budget.materialized.get().checked_sub(self.bytes) else {
610            self.budget.refuse_limit(
611                ResourceDimension::MaterializedBytes,
612                ResourceFailure::BudgetExceeded,
613                self.budget.materialized_allowance(),
614                self.budget.materialized.get(),
615                self.bytes,
616                self.operation,
617            );
618            return;
619        };
620        self.budget.materialized.set(remaining);
621    }
622}
623
624/// A guard holding one unit of recursive nesting depth.
625#[derive(Debug)]
626pub struct DepthGuard<'a> {
627    budget: &'a DecodeBudget,
628}
629
630impl Drop for DepthGuard<'_> {
631    fn drop(&mut self) {
632        let Some(depth) = self.budget.recursion_depth.get().checked_sub(1) else {
633            return;
634        };
635        self.budget.recursion_depth.set(depth);
636    }
637}
638
639/// A sticky, context-free local work counter.
640#[derive(Debug)]
641pub struct WorkBudget<'a> {
642    limit: usize,
643    remaining: Cell<Option<usize>>,
644    recursion_depth: Cell<usize>,
645    session: Option<&'a DecodeBudget>,
646    session_work_scale: NonZeroU64,
647}
648
649/// RAII guard for one recursive geometry-evaluation frame.
650#[derive(Debug)]
651pub struct WorkBudgetRecursionGuard<'budget, 'session> {
652    account: WorkRecursionAccount<'budget, 'session>,
653}
654
655#[derive(Debug)]
656enum WorkRecursionAccount<'budget, 'session> {
657    Session { _guard: DepthGuard<'session> },
658    Independent(&'budget WorkBudget<'session>),
659}
660
661impl Drop for WorkBudgetRecursionGuard<'_, '_> {
662    fn drop(&mut self) {
663        if let WorkRecursionAccount::Independent(budget) = &self.account {
664            if let Some(depth) = budget.recursion_depth.get().checked_sub(1) {
665                budget.recursion_depth.set(depth);
666            }
667        }
668    }
669}
670
671impl WorkBudget<'static> {
672    /// Creates an independent local work slice.
673    pub const fn new(limit: usize) -> Self {
674        Self {
675            limit,
676            remaining: Cell::new(Some(limit)),
677            recursion_depth: Cell::new(0),
678            session: None,
679            session_work_scale: NonZeroU64::MIN,
680        }
681    }
682}
683
684impl<'a> WorkBudget<'a> {
685    pub(super) fn for_session(limit: u64, session: &'a DecodeBudget) -> Self {
686        let Ok(limit) = usize::try_from(limit) else {
687            drop(session.refuse(
688                ResourceDimension::WorkUnits,
689                ResourceFailure::BudgetExceeded,
690                u64_from_index(usize::MAX),
691                0,
692                limit,
693                "work_budget",
694            ));
695            return Self {
696                limit: 0,
697                remaining: Cell::new(None),
698                recursion_depth: Cell::new(0),
699                session: Some(session),
700                session_work_scale: NonZeroU64::MIN,
701            };
702        };
703        Self {
704            limit,
705            remaining: Cell::new(Some(limit)),
706            recursion_depth: Cell::new(0),
707            session: Some(session),
708            session_work_scale: NonZeroU64::MIN,
709        }
710    }
711
712    /// Set the session work owed by each local unit without changing this
713    /// slice's local ceiling. Attached children preserve the same scale.
714    #[must_use]
715    pub fn with_session_work_scale(mut self, scale: NonZeroU64) -> Self {
716        self.session_work_scale = scale;
717        self
718    }
719
720    /// Reserve temporary evaluator storage in the attached decode session.
721    /// Independent slices retain fallible allocation without session charging.
722    pub fn reserve_scratch(
723        &self,
724        bytes: u64,
725        operation: &'static str,
726    ) -> Result<super::work_scratch::WorkScratch<'a>, ResourceLimit> {
727        super::work_scratch::WorkScratch::new(self.session, bytes, operation)
728    }
729
730    /// Charges one work unit, returning false after exhaustion.
731    pub fn charge(&self) -> bool {
732        self.charge_by(1)
733    }
734
735    /// Charges several work units, with sticky exhaustion on refusal.
736    pub fn charge_by(&self, work: usize) -> bool {
737        self.charge_by_against(work, self.session)
738    }
739
740    fn charge_by_against(&self, work: usize, session: Option<&DecodeBudget>) -> bool {
741        let Some(remaining) = self.remaining.get() else {
742            return false;
743        };
744        if work > remaining {
745            self.remaining.set(None);
746            false
747        } else {
748            if let Some(session) = session {
749                let Some(scaled) = u64_from_index(work).checked_mul(self.session_work_scale.get())
750                else {
751                    // The sticky session keeps the refusal for finish_session.
752                    let _failure = session.refuse(
753                        ResourceDimension::Codec("scaled_work_budget"),
754                        ResourceFailure::BudgetExceeded,
755                        u64::MAX - 1,
756                        u64::MAX - 1,
757                        1,
758                        "scaled_work_budget",
759                    );
760                    self.remaining.set(None);
761                    return false;
762                };
763                if session.charge_work(scaled, "work_budget").is_err() {
764                    self.remaining.set(None);
765                    return false;
766                }
767            }
768            self.remaining.set(Some(remaining - work));
769            true
770        }
771    }
772
773    /// Returns whether this slice has refused a charge.
774    pub fn exhausted(&self) -> bool {
775        self.remaining.get().is_none()
776    }
777
778    /// Marks this slice exhausted without consuming additional session work.
779    pub fn exhaust(&self) {
780        self.remaining.set(None);
781    }
782
783    /// Returns unconsumed work units.
784    pub fn remaining(&self) -> usize {
785        self.remaining.get().unwrap_or(0)
786    }
787
788    /// Returns charged work and the remainder forfeited by exhaustion.
789    pub fn consumed(&self) -> usize {
790        let remaining = self.remaining();
791        if remaining > self.limit {
792            self.exhaust();
793            return self.limit;
794        }
795        self.limit - remaining
796    }
797
798    /// Enters a session frame, or an independent frame with a 256-frame ceiling.
799    /// Attached child slices share the active session depth.
800    pub fn recursion_guard(&self) -> Result<WorkBudgetRecursionGuard<'_, 'a>, ResourceLimit> {
801        const INDEPENDENT_RECURSION_DEPTH: usize = 256;
802        if let Some(session) = self.session {
803            return session.enter_nested("work_budget_recursion").map(|guard| {
804                WorkBudgetRecursionGuard {
805                    account: WorkRecursionAccount::Session { _guard: guard },
806                }
807            });
808        }
809        let depth = self.recursion_depth.get();
810        if depth >= INDEPENDENT_RECURSION_DEPTH {
811            self.exhaust();
812            return Err(ResourceLimit {
813                dimension: ResourceDimension::RecursionDepth,
814                reason: ResourceFailure::BudgetExceeded,
815                limit: u64_from_index(INDEPENDENT_RECURSION_DEPTH),
816                used: u64_from_index(depth),
817                additional: 1,
818                operation: "work_budget_recursion",
819            });
820        }
821        self.recursion_depth.set(depth + 1);
822        Ok(WorkBudgetRecursionGuard {
823            account: WorkRecursionAccount::Independent(self),
824        })
825    }
826
827    /// Copy the active cycle path and admit one new frame slot.
828    /// The reservation covers the copied path until the frame leaves.
829    pub fn copy_recursion_path<T: Copy>(
830        &self,
831        path: &[T],
832    ) -> Result<(Vec<T>, super::work_scratch::WorkScratch<'a>), ResourceLimit> {
833        const OPERATION: &str = "model evaluation cycle path";
834        let capacity = path.len().checked_add(1).ok_or_else(|| ResourceLimit {
835            dimension: ResourceDimension::CollectionItems,
836            reason: ResourceFailure::BudgetExceeded,
837            limit: u64::MAX,
838            used: u64_from_index(path.len()),
839            additional: 1,
840            operation: OPERATION,
841        })?;
842        let bytes = u64_from_index(capacity)
843            .checked_mul(u64_from_index(std::mem::size_of::<T>()))
844            .ok_or(ResourceLimit {
845                dimension: ResourceDimension::MaterializedBytes,
846                reason: ResourceFailure::BudgetExceeded,
847                limit: u64::MAX,
848                used: 0,
849                additional: u64::MAX,
850                operation: OPERATION,
851            })?;
852        if let Some(session) = self.session {
853            session.charge_collection_items_limit(u64_from_index(capacity), OPERATION)?;
854            // Copy each path member and compare it once when binding the frame.
855            session.charge_work_limit(
856                bytes
857                    .checked_add(u64_from_index(path.len()))
858                    .ok_or_else(|| {
859                        session.refuse_limit(
860                            ResourceDimension::WorkUnits,
861                            ResourceFailure::BudgetExceeded,
862                            u64::MAX,
863                            bytes,
864                            u64_from_index(path.len()),
865                            OPERATION,
866                        )
867                    })?,
868                OPERATION,
869            )?;
870        }
871        let storage = self.reserve_scratch(bytes, OPERATION)?;
872        let mut copied = Vec::new();
873        copied
874            .try_reserve_exact(capacity)
875            .map_err(|_| match self.session {
876                Some(session) => session.refuse_limit(
877                    ResourceDimension::MaterializedBytes,
878                    ResourceFailure::AllocationFailed,
879                    session.materialized_allowance(),
880                    session.materialized.get(),
881                    bytes,
882                    OPERATION,
883                ),
884                None => ResourceLimit::allocation_failed(
885                    ResourceDimension::MaterializedBytes,
886                    bytes,
887                    bytes,
888                    OPERATION,
889                ),
890            })?;
891        copied.extend_from_slice(path);
892        Ok((copied, storage))
893    }
894
895    /// Creates an independent child slice capped by this budget's remainder.
896    pub fn child_slice(&self, limit: usize) -> WorkBudget<'static> {
897        WorkBudget::new(limit.min(self.remaining()))
898    }
899
900    /// Creates a child slice capped by this budget's remainder and attached to its session.
901    pub fn session_child_slice(&self, limit: usize) -> WorkBudget<'_> {
902        WorkBudget {
903            limit: limit.min(self.remaining()),
904            remaining: Cell::new(Some(limit.min(self.remaining()))),
905            recursion_depth: Cell::new(0),
906            session: self.session,
907            session_work_scale: self.session_work_scale,
908        }
909    }
910
911    /// Charges this budget for work consumed by a child slice.
912    ///
913    /// # Errors
914    ///
915    /// [`BudgetExhausted`] when the child's work is above this budget's
916    /// remainder. The budget marks itself exhausted in that case, so
917    /// [`WorkBudget::exhausted`] answers `true` afterwards.
918    pub fn consume_child(&self, child: &WorkBudget<'_>) -> Result<(), BudgetExhausted> {
919        // A session child already charged the shared session for each unit.
920        // Only transfer its consumption into this parent's local slice.
921        let session = match (self.session, child.session) {
922            (Some(parent), Some(child)) if std::ptr::eq(parent, child) => None,
923            _ => self.session,
924        };
925        if self.charge_by_against(child.consumed(), session) {
926            Ok(())
927        } else {
928            Err(BudgetExhausted)
929        }
930    }
931}
932
933/// A work budget reached its limit.
934///
935/// The budget marks itself exhausted when it refuses, so this error names the
936/// event and carries no state of its own.
937#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
938#[error("work budget exhausted")]
939pub struct BudgetExhausted;
940
941/// The work a step over `items` owes a [`WorkBudget`].
942///
943/// A step that walks no item still runs, so it owes one unit. Charging zero
944/// would leave an unbounded number of empty steps free of the budget, and the
945/// walk that makes them would not be bounded by it.
946#[must_use]
947pub const fn work_units(items: usize) -> usize {
948    match items {
949        0 => 1,
950        items => items,
951    }
952}
953
954/// Builds a correctly classified refusal for a codec-local ceiling.
955pub fn refuse_local_limit(what: &'static str, limit: u64, requested: u64) -> CodecError {
956    local_limit_error(what, limit, requested, what)
957}
958
959fn local_limit_error(
960    what: &'static str,
961    limit: u64,
962    requested: u64,
963    operation: &'static str,
964) -> CodecError {
965    let (used, additional) = match requested.checked_sub(limit) {
966        Some(excess) => (limit, excess),
967        None => (requested, 0),
968    };
969    CodecError::ResourceLimit(ResourceLimit {
970        dimension: ResourceDimension::Codec(what),
971        reason: ResourceFailure::BudgetExceeded,
972        limit,
973        used,
974        additional,
975        operation,
976    })
977}
978
979#[cfg(test)]
980mod tests {
981    use super::{work_units, DecodeBudget, WorkBudget};
982    use crate::decode::{DecodePolicy, ResourceDimension, ResourceFailure};
983
984    fn descend(budget: &WorkBudget<'_>, depth: usize) -> usize {
985        let Ok(_guard) = budget.recursion_guard() else {
986            return depth;
987        };
988        descend(budget, depth + 1)
989    }
990
991    #[test]
992    fn recursion_guard_is_shared_and_bounded() {
993        let budget = WorkBudget::new(10_000);
994        assert_eq!(descend(&budget, 0), 256);
995        assert!(budget.exhausted());
996    }
997
998    #[test]
999    fn attached_recursion_refuses_zero_session_depth() {
1000        let mut policy = DecodePolicy::default();
1001        policy.limits.max_recursion_depth = 0;
1002        let session = DecodeBudget::new(policy, 1);
1003        let budget = WorkBudget::for_session(100, &session);
1004        let failure = budget
1005            .recursion_guard()
1006            .expect_err("zero depth refuses first frame");
1007        assert_eq!(failure.dimension, ResourceDimension::RecursionDepth);
1008        assert_eq!(failure.limit, 0);
1009        assert_eq!(failure.used, 0);
1010        assert_eq!(failure.additional, 1);
1011        assert_eq!(session.fused(), Some(failure));
1012    }
1013
1014    #[test]
1015    fn attached_child_preserves_active_session_depth() {
1016        let mut policy = DecodePolicy::default();
1017        policy.limits.max_recursion_depth = 1;
1018        let session = DecodeBudget::new(policy, 1);
1019        let budget = WorkBudget::for_session(100, &session);
1020        let guard = budget.recursion_guard().expect("first frame fits");
1021        let child = budget.session_child_slice(100);
1022        let failure = child.recursion_guard().expect_err("child shares depth");
1023        assert_eq!(failure.dimension, ResourceDimension::RecursionDepth);
1024        assert_eq!(failure.used, 1);
1025        drop(guard);
1026        assert_eq!(session.recursion_depth.get(), 0);
1027    }
1028
1029    #[test]
1030    fn a_step_over_no_item_owes_exactly_one_unit() {
1031        assert_eq!(work_units(0), 1);
1032        for items in [1, 2, 7, usize::MAX] {
1033            assert_eq!(work_units(items), items);
1034        }
1035
1036        let budget = WorkBudget::new(1);
1037        assert!(budget.charge_by(work_units(0)));
1038        assert_eq!(budget.consumed(), 1);
1039        assert!(!budget.charge_by(work_units(0)));
1040        assert!(budget.exhausted());
1041    }
1042
1043    #[test]
1044    fn session_child_consumption_charges_the_session_once() {
1045        let mut policy = DecodePolicy::default();
1046        policy.limits.max_work_units = 1;
1047        let session = DecodeBudget::new(policy, 1);
1048        let parent = WorkBudget::for_session(2, &session);
1049        let child = parent.session_child_slice(1);
1050        assert!(child.charge());
1051        assert_eq!(session.work.get(), 1);
1052        assert!(parent.consume_child(&child).is_ok());
1053        assert_eq!(parent.remaining(), 1);
1054        assert_eq!(session.work.get(), 1);
1055        assert!(session.fused().is_none());
1056    }
1057
1058    #[test]
1059    fn scaled_session_work_preserves_local_ceiling_and_child_accounting() {
1060        let mut policy = DecodePolicy::service();
1061        policy.limits.max_work_units = 8;
1062        let session = DecodeBudget::new(policy, 1);
1063        let parent = WorkBudget::for_session(3, &session)
1064            .with_session_work_scale(std::num::NonZeroU64::new(4).expect("nonzero work scale"));
1065        let child = parent.session_child_slice(1);
1066        assert!(child.charge());
1067        assert_eq!(session.work.get(), 4);
1068        assert!(parent.consume_child(&child).is_ok());
1069        assert_eq!(parent.remaining(), 2);
1070        assert_eq!(session.work.get(), 4);
1071        assert!(parent.charge());
1072        assert_eq!(parent.remaining(), 1);
1073        assert_eq!(session.work.get(), 8);
1074        assert!(!parent.charge());
1075        assert!(parent.exhausted());
1076        let failure = session.fused().expect("session work refusal");
1077        assert_eq!(failure.dimension, super::ResourceDimension::WorkUnits);
1078        assert_eq!(failure.used, 8);
1079        assert_eq!(failure.additional, 4);
1080    }
1081
1082    #[test]
1083    fn retained_allocation_refusal_uses_the_retained_dimension_and_prior_usage() {
1084        let session = DecodeBudget::new(DecodePolicy::default(), 1);
1085        session
1086            .charge_retained(3, "copy retained")
1087            .expect("retained charge fits the default session allowance");
1088        let crate::CodecError::ResourceLimit(limit) =
1089            session.retained_allocation_failed(3, "copy retained")
1090        else {
1091            panic!("retained allocation must produce a resource refusal");
1092        };
1093        assert_eq!(limit.dimension, ResourceDimension::RetainedBytes);
1094        assert_eq!(limit.reason, ResourceFailure::AllocationFailed);
1095        assert_eq!(limit.used, 0);
1096        assert_eq!(limit.additional, 3);
1097        assert_eq!(session.fused(), Some(limit));
1098    }
1099
1100    #[test]
1101    fn zero_charge_refuses_when_usage_already_exceeds_limit() {
1102        let mut policy = DecodePolicy::default();
1103        policy.limits.max_entities = 1;
1104        let session = DecodeBudget::new(policy, 1);
1105        session.entities.set(2);
1106        let error = session.charge_entities(0, "entities");
1107        assert!(matches!(
1108            error,
1109            Err(crate::CodecError::ResourceLimit(limit))
1110                if limit.dimension == ResourceDimension::Entities
1111                    && limit.reason == ResourceFailure::BudgetExceeded
1112                    && limit.used == 2
1113                    && limit.additional == 0
1114        ));
1115    }
1116
1117    #[test]
1118    fn retained_allocation_refusal_reports_existing_usage_when_charge_exceeds_it() {
1119        let session = DecodeBudget::new(DecodePolicy::default(), 1);
1120        session.retained.set(2);
1121        let limit = session.retained_allocation_failed_limit(3, "retain");
1122        assert_eq!(limit.used, 2);
1123        assert_eq!(limit.additional, 3);
1124    }
1125}
cadmpeg_core/decode/budget.rs

cadmpeg_core/decode/
budget.rs
