hawdb_executor/
memory_ledger.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
15//! Query-owned hierarchical memory accounting.
16
17use hawdb_core::{HawDBError, Result};
18use std::collections::BTreeMap;
19use std::num::NonZeroUsize;
20use std::sync::{Arc, Mutex, MutexGuard};
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
23pub enum QueryMemoryClass {
24    PipelineBatch,
25    BlockingState,
26    ExternalRead,
27    SpillStaging,
28    MorselOutput,
29    ResultMaterialization,
30}
31
32impl QueryMemoryClass {
33    pub const fn as_str(self) -> &'static str {
34        match self {
35            Self::PipelineBatch => "pipeline_batch",
36            Self::BlockingState => "blocking_state",
37            Self::ExternalRead => "external_read",
38            Self::SpillStaging => "spill_staging",
39            Self::MorselOutput => "morsel_output",
40            Self::ResultMaterialization => "result_materialization",
41        }
42    }
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct QueryMemoryClassSnapshot {
47    pub class: QueryMemoryClass,
48    pub used_bytes: usize,
49    pub peak_bytes: usize,
50}
51
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub struct QueryMemoryLedgerSnapshot {
54    pub budget_bytes: usize,
55    pub used_bytes: usize,
56    pub peak_bytes: usize,
57    pub account_count: usize,
58    pub classes: Vec<QueryMemoryClassSnapshot>,
59}
60
61#[derive(Debug, Clone)]
62pub struct QueryMemoryLedger {
63    inner: Arc<QueryMemoryLedgerInner>,
64    file_descriptors: Option<hawdb_storage::file_descriptors::FileOpenContext>,
65}
66
67#[derive(Debug)]
68struct QueryMemoryLedgerInner {
69    budget_bytes: usize,
70    state: Mutex<QueryMemoryLedgerState>,
71}
72
73#[derive(Debug, Default)]
74struct QueryMemoryLedgerState {
75    used_bytes: usize,
76    peak_bytes: usize,
77    next_account_id: u64,
78    accounts: BTreeMap<u64, QueryMemoryAccountState>,
79    classes: BTreeMap<QueryMemoryClass, QueryMemoryClassState>,
80}
81
82#[derive(Debug)]
83struct QueryMemoryAccountState {
84    class: QueryMemoryClass,
85    owner: Arc<str>,
86    budget_bytes: usize,
87    used_bytes: usize,
88    peak_bytes: usize,
89}
90
91#[derive(Debug, Default)]
92struct QueryMemoryClassState {
93    used_bytes: usize,
94    peak_bytes: usize,
95}
96
97impl QueryMemoryLedger {
98    pub(crate) fn budget_bytes(&self) -> NonZeroUsize {
99        NonZeroUsize::new(self.inner.budget_bytes).unwrap()
100    }
101
102    pub(crate) fn source_account(
103        &self,
104        owner: &'static str,
105        source_budget: NonZeroUsize,
106        retained_budget: NonZeroUsize,
107    ) -> QueryMemoryAccount {
108        self.source_account_retaining(
109            owner,
110            source_budget,
111            self.account(QueryMemoryClass::BlockingState, owner, retained_budget),
112        )
113    }
114
115    pub(crate) fn source_account_retaining(
116        &self,
117        owner: &'static str,
118        source_budget: NonZeroUsize,
119        retained: QueryMemoryAccount,
120    ) -> QueryMemoryAccount {
121        self.account(QueryMemoryClass::ExternalRead, owner, source_budget)
122            .with_retained_state(retained)
123    }
124
125    pub fn new(budget_bytes: NonZeroUsize) -> Self {
126        Self {
127            inner: Arc::new(QueryMemoryLedgerInner {
128                budget_bytes: budget_bytes.get(),
129                state: Mutex::new(QueryMemoryLedgerState::default()),
130            }),
131            file_descriptors: None,
132        }
133    }
134
135    /// Bind query IO to its store while retaining the existing memory ledger.
136    #[doc(hidden)]
137    pub fn with_file_descriptors(
138        mut self,
139        context: Option<hawdb_storage::file_descriptors::FileOpenContext>,
140    ) -> Self {
141        self.file_descriptors = context.or(self.file_descriptors);
142        self
143    }
144
145    pub(crate) fn file_descriptor_context(
146        &self,
147    ) -> Option<&hawdb_storage::file_descriptors::FileOpenContext> {
148        self.file_descriptors.as_ref()
149    }
150
151    pub fn account(
152        &self,
153        class: QueryMemoryClass,
154        owner: impl Into<Arc<str>>,
155        budget_bytes: NonZeroUsize,
156    ) -> QueryMemoryAccount {
157        // Caller code must run before locking: conversion can panic or re-enter the ledger.
158        let owner = owner.into();
159        let mut state = lock_recover(&self.inner.state);
160        let account_id = state.next_account_id;
161        state.next_account_id = state.next_account_id.saturating_add(1);
162        state.accounts.insert(
163            account_id,
164            QueryMemoryAccountState {
165                class,
166                owner,
167                budget_bytes: budget_bytes.get(),
168                used_bytes: 0,
169                peak_bytes: 0,
170            },
171        );
172        QueryMemoryAccount {
173            ledger: self.clone(),
174            account_id,
175            backing: None,
176            retained_state: None,
177        }
178    }
179
180    pub fn snapshot(&self) -> QueryMemoryLedgerSnapshot {
181        let state = lock_recover(&self.inner.state);
182        QueryMemoryLedgerSnapshot {
183            budget_bytes: self.inner.budget_bytes,
184            used_bytes: state.used_bytes,
185            peak_bytes: state.peak_bytes,
186            account_count: state.accounts.len(),
187            classes: state
188                .classes
189                .iter()
190                .map(|(class, state)| QueryMemoryClassSnapshot {
191                    class: *class,
192                    used_bytes: state.used_bytes,
193                    peak_bytes: state.peak_bytes,
194                })
195                .collect(),
196        }
197    }
198
199    /// Test observation of the actual owner, rather than unrelated root charges.
200    #[cfg(test)]
201    pub(crate) fn test_owner_used_bytes(&self, class: QueryMemoryClass, owner: &str) -> usize {
202        lock_recover(&self.inner.state)
203            .accounts
204            .values()
205            .filter(|account| account.class == class && account.owner.as_ref() == owner)
206            .map(|account| account.used_bytes)
207            .sum()
208    }
209
210    fn reserve(&self, account_id: u64, bytes: usize) -> Result<()> {
211        if bytes == 0 {
212            return Ok(());
213        }
214        let mut state = lock_recover(&self.inner.state);
215        let (class, owner, account_budget, account_next) = {
216            let account = state.accounts.get(&account_id).ok_or_else(|| {
217                HawDBError::Execution("query memory account is no longer registered".to_string())
218            })?;
219            let account_next = account.used_bytes.checked_add(bytes).ok_or_else(|| {
220                HawDBError::Execution(format!(
221                    "query memory account {} ({}) byte accounting overflow",
222                    account.owner,
223                    account.class.as_str()
224                ))
225            })?;
226            (
227                account.class,
228                Arc::clone(&account.owner),
229                account.budget_bytes,
230                account_next,
231            )
232        };
233        if account_next > account_budget {
234            return Err(HawDBError::Execution(format!(
235                "query memory account {owner} ({}) would use {account_next} bytes, exceeding its {account_budget}-byte budget",
236                class.as_str()
237            )));
238        }
239        let root_next = state.used_bytes.checked_add(bytes).ok_or_else(|| {
240            HawDBError::Execution(format!(
241                "query memory ledger byte accounting overflow while charging {owner} ({})",
242                class.as_str()
243            ))
244        })?;
245        if root_next > self.inner.budget_bytes {
246            return Err(HawDBError::Execution(format!(
247                "query memory ledger would use {root_next} bytes while charging {owner} ({}), exceeding query_memory_bytes {}",
248                class.as_str(),
249                self.inner.budget_bytes
250            )));
251        }
252
253        // Complete allocation before changing any of the hierarchical counters.
254        state.classes.entry(class).or_default();
255        let account = state
256            .accounts
257            .get_mut(&account_id)
258            .expect("validated query memory account remains registered");
259        account.used_bytes = account_next;
260        account.peak_bytes = account.peak_bytes.max(account_next);
261        state.used_bytes = root_next;
262        state.peak_bytes = state.peak_bytes.max(root_next);
263        let class_state = state
264            .classes
265            .get_mut(&class)
266            .expect("prepared query memory class remains registered");
267        class_state.used_bytes = class_state.used_bytes.saturating_add(bytes);
268        class_state.peak_bytes = class_state.peak_bytes.max(class_state.used_bytes);
269        Ok(())
270    }
271
272    fn can_reserve(&self, account_id: u64, bytes: usize) -> bool {
273        let state = lock_recover(&self.inner.state);
274        let Some(account) = state.accounts.get(&account_id) else {
275            return false;
276        };
277        let Some(account_next) = account.used_bytes.checked_add(bytes) else {
278            return false;
279        };
280        let Some(root_next) = state.used_bytes.checked_add(bytes) else {
281            return false;
282        };
283        account_next <= account.budget_bytes && root_next <= self.inner.budget_bytes
284    }
285
286    fn release(&self, account_id: u64, bytes: usize) {
287        if bytes == 0 {
288            return;
289        }
290        let mut state = lock_recover(&self.inner.state);
291        let Some(account) = state.accounts.get_mut(&account_id) else {
292            return;
293        };
294        let released = bytes.min(account.used_bytes);
295        let class = account.class;
296        account.used_bytes -= released;
297        state.used_bytes = state.used_bytes.saturating_sub(released);
298        if let Some(class_state) = state.classes.get_mut(&class) {
299            class_state.used_bytes = class_state.used_bytes.saturating_sub(released);
300        }
301    }
302
303    fn transfer(
304        &self,
305        source_account_id: u64,
306        source_bytes: usize,
307        target_account_id: u64,
308        target_bytes: usize,
309    ) -> Result<()> {
310        if source_account_id == target_account_id {
311            return Err(HawDBError::Execution(
312                "query memory transfer requires distinct accounts".to_string(),
313            ));
314        }
315        let mut state = lock_recover(&self.inner.state);
316        let (source_class, source_used) = state
317            .accounts
318            .get(&source_account_id)
319            .map(|account| (account.class, account.used_bytes))
320            .ok_or_else(|| {
321                HawDBError::Execution(
322                    "source query memory account is no longer registered".to_string(),
323                )
324            })?;
325        if source_bytes > source_used {
326            return Err(HawDBError::Execution(format!(
327                "query memory transfer tried to release {source_bytes} bytes from a source account using {source_used} bytes"
328            )));
329        }
330        let (target_class, target_owner, target_budget, target_used) = state
331            .accounts
332            .get(&target_account_id)
333            .map(|account| {
334                (
335                    account.class,
336                    Arc::clone(&account.owner),
337                    account.budget_bytes,
338                    account.used_bytes,
339                )
340            })
341            .ok_or_else(|| {
342                HawDBError::Execution(
343                    "target query memory account is no longer registered".to_string(),
344                )
345            })?;
346        let target_next = target_used.checked_add(target_bytes).ok_or_else(|| {
347            HawDBError::Execution(format!(
348                "query memory account {target_owner} ({}) byte accounting overflow",
349                target_class.as_str()
350            ))
351        })?;
352        if target_next > target_budget {
353            return Err(HawDBError::Execution(format!(
354                "query memory account {target_owner} ({}) would use {target_next} bytes, exceeding its {target_budget}-byte budget",
355                target_class.as_str()
356            )));
357        }
358        let root_after_release = state.used_bytes.checked_sub(source_bytes).ok_or_else(|| {
359            HawDBError::Execution(
360                "query memory ledger underflow during ownership transfer".to_string(),
361            )
362        })?;
363        let root_next = root_after_release.checked_add(target_bytes).ok_or_else(|| {
364            HawDBError::Execution(format!(
365                "query memory ledger byte accounting overflow while transferring ownership to {target_owner} ({})",
366                target_class.as_str()
367            ))
368        })?;
369        if root_next > self.inner.budget_bytes {
370            return Err(HawDBError::Execution(format!(
371                "query memory ledger would use {root_next} bytes while transferring ownership to {target_owner} ({}), exceeding query_memory_bytes {}",
372                target_class.as_str(),
373                self.inner.budget_bytes
374            )));
375        }
376
377        state.classes.entry(target_class).or_default();
378        let source = state
379            .accounts
380            .get_mut(&source_account_id)
381            .expect("validated source query memory account remains registered");
382        source.used_bytes -= source_bytes;
383        let target = state
384            .accounts
385            .get_mut(&target_account_id)
386            .expect("validated target query memory account remains registered");
387        target.used_bytes = target_next;
388        target.peak_bytes = target.peak_bytes.max(target_next);
389        state.used_bytes = root_next;
390        state.peak_bytes = state.peak_bytes.max(root_next);
391        if let Some(class_state) = state.classes.get_mut(&source_class) {
392            class_state.used_bytes = class_state.used_bytes.saturating_sub(source_bytes);
393        }
394        let class_state = state
395            .classes
396            .get_mut(&target_class)
397            .expect("prepared target query memory class remains registered");
398        class_state.used_bytes = class_state.used_bytes.saturating_add(target_bytes);
399        class_state.peak_bytes = class_state.peak_bytes.max(class_state.used_bytes);
400        Ok(())
401    }
402}
403
404#[derive(Debug, Clone)]
405pub struct QueryMemoryAccount {
406    ledger: QueryMemoryLedger,
407    account_id: u64,
408    // A child ledger owns an up-front reservation in its parent. Every account
409    // and outstanding lease keeps that reservation alive.
410    backing: Option<Arc<QueryMemoryLease>>,
411    retained_state: Option<Arc<QueryMemoryAccount>>,
412}
413
414impl QueryMemoryAccount {
415    pub(crate) fn query_budget_bytes(&self) -> NonZeroUsize {
416        self.ledger.budget_bytes()
417    }
418
419    pub(crate) fn with_retained_state(mut self, account: Self) -> Self {
420        self.retained_state = Some(Arc::new(account));
421        self
422    }
423
424    pub(crate) fn retained_state(&self) -> Option<&Self> {
425        self.retained_state.as_deref()
426    }
427
428    /// Reserve the entire child allowance before dispatch, without charging
429    /// individual child allocations to the query root a second time.
430    pub(crate) fn sub_account(&self, budget: NonZeroUsize) -> Result<Self> {
431        let backing = Arc::new(self.reserve(budget.get())?);
432        let (class, owner) = {
433            let state = lock_recover(&self.ledger.inner.state);
434            let account = &state.accounts[&self.account_id];
435            (account.class, Arc::clone(&account.owner))
436        };
437        let mut child = QueryMemoryLedger::new(budget).account(class, owner, budget);
438        child.backing = Some(backing);
439        Ok(child)
440    }
441
442    pub(crate) fn budget_bytes(&self) -> NonZeroUsize {
443        NonZeroUsize::new(
444            lock_recover(&self.ledger.inner.state).accounts[&self.account_id].budget_bytes,
445        )
446        .unwrap()
447    }
448
449    pub(crate) fn available_bytes(&self) -> usize {
450        let state = lock_recover(&self.ledger.inner.state);
451        let account = &state.accounts[&self.account_id];
452        account.budget_bytes.saturating_sub(account.used_bytes).min(
453            self.ledger
454                .inner
455                .budget_bytes
456                .saturating_sub(state.used_bytes),
457        )
458    }
459
460    pub(crate) fn peak_bytes(&self) -> usize {
461        lock_recover(&self.ledger.inner.state).accounts[&self.account_id].peak_bytes
462    }
463
464    pub(crate) fn sibling(
465        &self,
466        class: QueryMemoryClass,
467        owner: impl Into<Arc<str>>,
468        budget_bytes: NonZeroUsize,
469    ) -> Self {
470        let mut account = self.ledger.account(class, owner, budget_bytes);
471        account.backing = self.backing.clone();
472        account
473    }
474
475    pub(crate) fn can_reserve(&self, bytes: usize) -> bool {
476        self.ledger.can_reserve(self.account_id, bytes)
477    }
478
479    pub fn reserve(&self, bytes: usize) -> Result<QueryMemoryLease> {
480        self.ledger.reserve(self.account_id, bytes)?;
481        Ok(QueryMemoryLease {
482            account: self.clone(),
483            bytes,
484        })
485    }
486}
487
488#[derive(Debug)]
489pub struct QueryMemoryLease {
490    account: QueryMemoryAccount,
491    bytes: usize,
492}
493
494impl QueryMemoryLease {
495    pub(crate) fn belongs_to(&self, account: &QueryMemoryAccount) -> bool {
496        self.account.account_id == account.account_id
497            && Arc::ptr_eq(&self.account.ledger.inner, &account.ledger.inner)
498    }
499
500    pub fn bytes(&self) -> usize {
501        self.bytes
502    }
503
504    pub fn grow(&mut self, bytes: usize) -> Result<()> {
505        self.account
506            .ledger
507            .reserve(self.account.account_id, bytes)?;
508        self.bytes = self.bytes.saturating_add(bytes);
509        Ok(())
510    }
511
512    pub fn shrink(&mut self, bytes: usize) {
513        let released = bytes.min(self.bytes);
514        self.account
515            .ledger
516            .release(self.account.account_id, released);
517        self.bytes -= released;
518    }
519
520    pub fn reset(&mut self) {
521        self.shrink(self.bytes);
522    }
523
524    pub(crate) fn transfer_to(
525        &mut self,
526        source_bytes: usize,
527        target: &mut Self,
528        target_bytes: usize,
529    ) -> Result<()> {
530        if !Arc::ptr_eq(&self.account.ledger.inner, &target.account.ledger.inner) {
531            return Err(HawDBError::Execution(
532                "query memory transfer requires accounts from the same ledger".to_string(),
533            ));
534        }
535        if source_bytes > self.bytes {
536            return Err(HawDBError::Execution(format!(
537                "query memory transfer tried to release {source_bytes} bytes from a {}-byte lease",
538                self.bytes,
539            )));
540        }
541        self.account.ledger.transfer(
542            self.account.account_id,
543            source_bytes,
544            target.account.account_id,
545            target_bytes,
546        )?;
547        self.bytes -= source_bytes;
548        target.bytes = target.bytes.saturating_add(target_bytes);
549        Ok(())
550    }
551}
552
553impl Drop for QueryMemoryLease {
554    fn drop(&mut self) {
555        self.account
556            .ledger
557            .release(self.account.account_id, self.bytes);
558        self.bytes = 0;
559    }
560}
561
562fn lock_recover<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
563    // State never escapes the ledger. Checks and allocation precede counter updates,
564    // and no caller code runs under this lock. Recovering poisoning preserves that
565    // valid state; it is not a repair mechanism for arbitrary accounting corruption.
566    mutex
567        .lock()
568        .unwrap_or_else(|poisoned| poisoned.into_inner())
569}
570
571#[cfg(test)]
572#[path = "memory_ledger_hardening_tests.rs"]
573mod hardening_tests;
574
575#[cfg(test)]
576mod tests {
577    use super::*;
578
579    #[test]
580    fn source_account_reuses_retained_cap_without_registering_an_unused_account() {
581        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(1024).unwrap());
582        let retained = ledger.account(
583            QueryMemoryClass::BlockingState,
584            "shared state",
585            NonZeroUsize::new(512).unwrap(),
586        );
587        let source =
588            ledger.source_account_retaining("source", ledger.budget_bytes(), retained.clone());
589        assert_eq!(ledger.snapshot().account_count, 2);
590        let mut grant = crate::store::admit_graph_read(&source, None, 512).unwrap();
591        grant.retain_state().unwrap();
592        assert!(retained.reserve(1).is_err());
593        assert_eq!(ledger.snapshot().used_bytes, 512);
594        drop(grant);
595        let reused = retained.reserve(512).unwrap();
596        assert_eq!(ledger.snapshot().account_count, 2);
597        drop(reused);
598        assert_eq!(ledger.snapshot().used_bytes, 0);
599    }
600
601    #[test]
602    fn sub_accounts_reserve_parent_capacity_without_double_charging() {
603        let budget = NonZeroUsize::new(100).unwrap();
604        let ledger = QueryMemoryLedger::new(budget);
605        let parent = ledger.account(QueryMemoryClass::BlockingState, "merge", budget);
606        let retained = parent.reserve(10).unwrap();
607        let left = parent.sub_account(NonZeroUsize::new(60).unwrap()).unwrap();
608        let right = parent.sub_account(NonZeroUsize::new(30).unwrap()).unwrap();
609        assert_eq!(parent.available_bytes(), 0);
610        assert!(parent.sub_account(NonZeroUsize::MIN).is_err());
611        let mut left_lease = left.reserve(60).unwrap();
612        let right_lease = right.reserve(30).unwrap();
613        assert!(left_lease.grow(1).is_err());
614        assert_eq!(ledger.snapshot().used_bytes, 100);
615        drop((left, right, retained));
616        assert_eq!(ledger.snapshot().used_bytes, 90);
617        drop(left_lease);
618        assert_eq!(ledger.snapshot().used_bytes, 30);
619        drop(right_lease);
620        assert_eq!(ledger.snapshot().used_bytes, 0);
621        assert_eq!(ledger.snapshot().peak_bytes, 100);
622    }
623
624    #[test]
625    fn nested_accounts_and_siblings_keep_backing_until_last_owner_drops() {
626        let budget = NonZeroUsize::new(8).unwrap();
627        let ledger = QueryMemoryLedger::new(budget);
628        let parent = ledger.account(QueryMemoryClass::BlockingState, "merge", budget);
629        let child = parent.sub_account(budget).unwrap();
630        let sibling = child.sibling(QueryMemoryClass::SpillStaging, "staging", budget);
631        let nested = child.sub_account(NonZeroUsize::new(4).unwrap()).unwrap();
632        assert!(sibling.reserve(5).is_err());
633        let lease = nested.reserve(4).unwrap();
634        drop((parent, child, nested));
635        assert_eq!(ledger.snapshot().used_bytes, 8);
636        drop(lease);
637        assert_eq!(ledger.snapshot().used_bytes, 8);
638        let sibling_lease = sibling.reserve(8).unwrap();
639        drop(sibling);
640        assert_eq!(ledger.snapshot().used_bytes, 8);
641        drop(sibling_lease);
642        assert_eq!(ledger.snapshot().used_bytes, 0);
643    }
644
645    #[test]
646    fn failed_sub_account_admission_does_not_leak_root_capacity() {
647        let budget = NonZeroUsize::new(100).unwrap();
648        let ledger = QueryMemoryLedger::new(budget);
649        let parent = ledger.account(QueryMemoryClass::BlockingState, "merge", budget);
650        let other = ledger.account(QueryMemoryClass::PipelineBatch, "input", budget);
651        let retained = other.reserve(70).unwrap();
652        assert!(parent.sub_account(NonZeroUsize::new(31).unwrap()).is_err());
653        assert_eq!(ledger.snapshot().used_bytes, 70);
654        let child = parent.sub_account(NonZeroUsize::new(30).unwrap()).unwrap();
655        assert!(child.reserve(31).is_err());
656        drop((child, retained));
657        assert_eq!(ledger.snapshot().used_bytes, 0);
658    }
659
660    #[test]
661    fn sibling_accounts_share_one_root_budget() {
662        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(10).unwrap());
663        let left = ledger.account(
664            QueryMemoryClass::BlockingState,
665            "left",
666            NonZeroUsize::new(10).unwrap(),
667        );
668        let right = ledger.account(
669            QueryMemoryClass::PipelineBatch,
670            "right",
671            NonZeroUsize::new(10).unwrap(),
672        );
673        let left_lease = left.reserve(6).unwrap();
674        let error = right.reserve(5).unwrap_err();
675
676        assert!(error.to_string().contains("query_memory_bytes 10"));
677        assert_eq!(ledger.snapshot().used_bytes, 6);
678        drop(left_lease);
679        assert_eq!(ledger.snapshot().used_bytes, 0);
680    }
681
682    #[test]
683    fn capacity_preflight_checks_account_and_root_without_charging() {
684        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(10).unwrap());
685        let left = ledger.account(
686            QueryMemoryClass::BlockingState,
687            "left",
688            NonZeroUsize::new(8).unwrap(),
689        );
690        let right = ledger.account(
691            QueryMemoryClass::PipelineBatch,
692            "right",
693            NonZeroUsize::new(10).unwrap(),
694        );
695
696        assert!(left.can_reserve(8));
697        assert_eq!(ledger.snapshot().used_bytes, 0);
698        let left_lease = left.reserve(6).unwrap();
699        assert!(!left.can_reserve(3));
700        assert!(right.can_reserve(4));
701        assert!(!right.can_reserve(5));
702        assert_eq!(ledger.snapshot().used_bytes, 6);
703
704        drop(left_lease);
705        assert_eq!(ledger.snapshot().used_bytes, 0);
706    }
707
708    #[test]
709    fn derived_sibling_account_keeps_the_original_query_root() {
710        let budget = NonZeroUsize::new(8).unwrap();
711        let ledger = QueryMemoryLedger::new(budget);
712        let state = ledger.account(QueryMemoryClass::BlockingState, "merge", budget);
713        let output = state.sibling(QueryMemoryClass::PipelineBatch, "output", budget);
714        let state_lease = state.reserve(6).unwrap();
715        assert!(output
716            .reserve(3)
717            .unwrap_err()
718            .to_string()
719            .contains("query_memory_bytes 8"));
720        let output_lease = output.reserve(2).unwrap();
721        assert_eq!(ledger.snapshot().used_bytes, 8);
722        drop((state_lease, output_lease));
723        assert_eq!(ledger.snapshot().used_bytes, 0);
724    }
725
726    #[test]
727    fn ownership_transfer_does_not_require_double_root_capacity() {
728        let budget = NonZeroUsize::new(8).unwrap();
729        let ledger = QueryMemoryLedger::new(budget);
730        let source = ledger.account(QueryMemoryClass::BlockingState, "source", budget);
731        let target = ledger.account(QueryMemoryClass::PipelineBatch, "target", budget);
732        let mut source_lease = source.reserve(8).unwrap();
733        let mut target_lease = target.reserve(0).unwrap();
734
735        source_lease.transfer_to(8, &mut target_lease, 8).unwrap();
736
737        assert_eq!(source_lease.bytes(), 0);
738        assert_eq!(target_lease.bytes(), 8);
739        let snapshot = ledger.snapshot();
740        assert_eq!(snapshot.used_bytes, 8);
741        assert_eq!(snapshot.peak_bytes, 8);
742        drop((source_lease, target_lease));
743        assert_eq!(ledger.snapshot().used_bytes, 0);
744    }
745
746    #[test]
747    fn failed_ownership_transfer_preserves_source_charge() {
748        let budget = NonZeroUsize::new(8).unwrap();
749        let ledger = QueryMemoryLedger::new(budget);
750        let source = ledger.account(QueryMemoryClass::BlockingState, "source", budget);
751        let target = ledger.account(QueryMemoryClass::PipelineBatch, "target", budget);
752        let mut source_lease = source.reserve(8).unwrap();
753        let mut target_lease = target.reserve(0).unwrap();
754
755        let error = source_lease
756            .transfer_to(8, &mut target_lease, 9)
757            .unwrap_err();
758
759        assert!(error.to_string().contains("9 bytes"));
760        assert_eq!(source_lease.bytes(), 8);
761        assert_eq!(target_lease.bytes(), 0);
762        assert_eq!(ledger.snapshot().used_bytes, 8);
763    }
764
765    #[test]
766    fn account_budget_is_enforced_before_root_budget() {
767        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(100).unwrap());
768        let account = ledger.account(
769            QueryMemoryClass::BlockingState,
770            "sort",
771            NonZeroUsize::new(4).unwrap(),
772        );
773
774        let error = account.reserve(5).unwrap_err();
775
776        assert!(error.to_string().contains("sort"));
777        assert!(error.to_string().contains("4-byte budget"));
778        assert_eq!(ledger.snapshot().used_bytes, 0);
779    }
780
781    #[test]
782    fn lease_growth_shrink_and_drop_update_hierarchy() {
783        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(100).unwrap());
784        let account = ledger.account(
785            QueryMemoryClass::MorselOutput,
786            "reorder",
787            NonZeroUsize::new(80).unwrap(),
788        );
789        let mut lease = account.reserve(20).unwrap();
790        lease.grow(30).unwrap();
791        lease.shrink(10);
792        assert_eq!(lease.bytes(), 40);
793        assert_eq!(ledger.snapshot().used_bytes, 40);
794
795        drop(lease);
796        let snapshot = ledger.snapshot();
797        assert_eq!(snapshot.used_bytes, 0);
798        assert_eq!(snapshot.peak_bytes, 50);
799        assert_eq!(snapshot.classes[0].peak_bytes, 50);
800    }
801
802    #[test]
803    fn unwind_releases_query_memory_lease() {
804        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(100).unwrap());
805        let account = ledger.account(
806            QueryMemoryClass::SpillStaging,
807            "spill",
808            NonZeroUsize::new(100).unwrap(),
809        );
810        let result = std::panic::catch_unwind(|| {
811            let _lease = account.reserve(80).unwrap();
812            panic!("injected panic");
813        });
814
815        assert!(result.is_err());
816        assert_eq!(ledger.snapshot().used_bytes, 0);
817    }
818}
hawdb_executor/memory_ledger.rs

hawdb_executor/
memory_ledger.rs
