hawdb_executor/

memory_ledger.rs

// Copyright 2026 Nowledge
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Query-owned hierarchical memory accounting.

use hawdb_core::{
    HawDBError, ResourceBudgetError, ResourceBudgetKind, ResourceBudgetScope, Result,
};
use std::collections::BTreeMap;
use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex, MutexGuard};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum QueryMemoryClass {
    PipelineBatch,
    BlockingState,
    ExternalRead,
    SpillStaging,
    MorselOutput,
    ResultMaterialization,
}

impl QueryMemoryClass {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::PipelineBatch => "pipeline_batch",
            Self::BlockingState => "blocking_state",
            Self::ExternalRead => "external_read",
            Self::SpillStaging => "spill_staging",
            Self::MorselOutput => "morsel_output",
            Self::ResultMaterialization => "result_materialization",
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QueryMemoryClassSnapshot {
    pub class: QueryMemoryClass,
    pub used_bytes: usize,
    pub peak_bytes: usize,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QueryMemoryLedgerSnapshot {
    pub budget_bytes: usize,
    pub used_bytes: usize,
    pub peak_bytes: usize,
    pub account_count: usize,
    pub classes: Vec<QueryMemoryClassSnapshot>,
}

#[derive(Debug, Clone)]
pub struct QueryMemoryLedger {
    inner: Arc<QueryMemoryLedgerInner>,
    file_descriptors: Option<hawdb_storage::file_descriptors::FileOpenContext>,
}

#[derive(Debug)]
struct QueryMemoryLedgerInner {
    budget_bytes: usize,
    budget_option: &'static str,
    state: Mutex<QueryMemoryLedgerState>,
}

#[derive(Debug, Default)]
struct QueryMemoryLedgerState {
    used_bytes: usize,
    peak_bytes: usize,
    next_account_id: u64,
    accounts: BTreeMap<u64, QueryMemoryAccountState>,
    classes: BTreeMap<QueryMemoryClass, QueryMemoryClassState>,
}

#[derive(Debug)]
struct QueryMemoryAccountState {
    class: QueryMemoryClass,
    owner: Arc<str>,
    budget_bytes: usize,
    budget_option: &'static str,
    used_bytes: usize,
    peak_bytes: usize,
}

#[derive(Debug, Default)]
struct QueryMemoryClassState {
    used_bytes: usize,
    peak_bytes: usize,
}

impl QueryMemoryAccountState {
    fn budget_error(
        &self,
        scope: ResourceBudgetScope,
        requested: usize,
        used: usize,
        budget: usize,
        budget_option: &'static str,
        message: String,
    ) -> HawDBError {
        HawDBError::from(ResourceBudgetError {
            resource: ResourceBudgetKind::MemoryBytes,
            scope,
            owner: self.owner.to_string(),
            memory_class: Some(self.class.as_str()),
            requested: requested as u64,
            used: used as u64,
            budget: budget as u64,
            budget_option,
            message,
        })
    }

    fn account_budget_error(&self, requested: usize, message: String) -> HawDBError {
        self.budget_error(
            ResourceBudgetScope::Account,
            requested,
            self.used_bytes,
            self.budget_bytes,
            self.budget_option,
            message,
        )
    }
}

impl QueryMemoryLedger {
    /// Creates a query root with the effective host or runtime-granted limit.
    #[doc(hidden)]
    pub fn for_query(
        memory: &crate::ExecutionMemoryConfig,
        task_context: Option<&hawdb_core::RuntimeTaskContext>,
    ) -> Result<Self> {
        Ok(Self::new_with_budget_option(
            crate::memory::enforced_query_memory_budget(memory, task_context)?,
            if task_context
                .and_then(hawdb_core::RuntimeTaskContext::memory_reservation)
                .is_some()
            {
                "RuntimeTaskContext.memory_reservation.memory_bytes"
            } else {
                "ExecutionMemoryConfig.query_memory_bytes"
            },
        ))
    }

    pub(crate) fn budget_bytes(&self) -> NonZeroUsize {
        NonZeroUsize::new(self.inner.budget_bytes).unwrap()
    }

    pub(crate) fn source_account(
        &self,
        owner: &'static str,
        source_budget: NonZeroUsize,
        retained_budget: NonZeroUsize,
    ) -> QueryMemoryAccount {
        self.source_account_retaining(
            owner,
            source_budget,
            self.account(QueryMemoryClass::BlockingState, owner, retained_budget),
        )
    }

    pub(crate) fn source_account_retaining(
        &self,
        owner: &'static str,
        source_budget: NonZeroUsize,
        retained: QueryMemoryAccount,
    ) -> QueryMemoryAccount {
        self.account(QueryMemoryClass::ExternalRead, owner, source_budget)
            .with_retained_state(retained)
    }

    pub fn new(budget_bytes: NonZeroUsize) -> Self {
        Self::new_with_budget_option(budget_bytes, "QueryMemoryLedger::new(budget_bytes)")
    }

    /// Names the configuration option or runtime grant supplying this limit.
    pub fn new_with_budget_option(budget_bytes: NonZeroUsize, budget_option: &'static str) -> Self {
        Self {
            inner: Arc::new(QueryMemoryLedgerInner {
                budget_bytes: budget_bytes.get(),
                budget_option,
                state: Mutex::new(QueryMemoryLedgerState::default()),
            }),
            file_descriptors: None,
        }
    }

    /// Bind query IO to its store while retaining the existing memory ledger.
    #[doc(hidden)]
    pub fn with_file_descriptors(
        mut self,
        context: Option<hawdb_storage::file_descriptors::FileOpenContext>,
    ) -> Self {
        self.file_descriptors = context.or(self.file_descriptors);
        self
    }

    pub(crate) fn file_descriptor_context(
        &self,
    ) -> Option<&hawdb_storage::file_descriptors::FileOpenContext> {
        self.file_descriptors.as_ref()
    }

    pub fn account(
        &self,
        class: QueryMemoryClass,
        owner: impl Into<Arc<str>>,
        budget_bytes: NonZeroUsize,
    ) -> QueryMemoryAccount {
        self.account_with_budget_option(
            class,
            owner,
            budget_bytes,
            "QueryMemoryLedger::account(budget_bytes)",
        )
    }

    /// Creates an account with an explicit origin for its effective byte limit.
    pub fn account_with_budget_option(
        &self,
        class: QueryMemoryClass,
        owner: impl Into<Arc<str>>,
        budget_bytes: NonZeroUsize,
        budget_option: &'static str,
    ) -> QueryMemoryAccount {
        // Caller code must run before locking: conversion can panic or re-enter the ledger.
        let owner = owner.into();
        let mut state = lock_recover(&self.inner.state);
        let account_id = state.next_account_id;
        state.next_account_id = state.next_account_id.saturating_add(1);
        state.accounts.insert(
            account_id,
            QueryMemoryAccountState {
                class,
                owner,
                budget_bytes: budget_bytes.get(),
                budget_option,
                used_bytes: 0,
                peak_bytes: 0,
            },
        );
        QueryMemoryAccount {
            ledger: self.clone(),
            account_id,
            backing: None,
            retained_state: None,
        }
    }

    pub fn snapshot(&self) -> QueryMemoryLedgerSnapshot {
        let state = lock_recover(&self.inner.state);
        QueryMemoryLedgerSnapshot {
            budget_bytes: self.inner.budget_bytes,
            used_bytes: state.used_bytes,
            peak_bytes: state.peak_bytes,
            account_count: state.accounts.len(),
            classes: state
                .classes
                .iter()
                .map(|(class, state)| QueryMemoryClassSnapshot {
                    class: *class,
                    used_bytes: state.used_bytes,
                    peak_bytes: state.peak_bytes,
                })
                .collect(),
        }
    }

    /// Test observation of the actual owner, rather than unrelated root charges.
    #[cfg(test)]
    pub(crate) fn test_owner_used_bytes(&self, class: QueryMemoryClass, owner: &str) -> usize {
        lock_recover(&self.inner.state)
            .accounts
            .values()
            .filter(|account| account.class == class && account.owner.as_ref() == owner)
            .map(|account| account.used_bytes)
            .sum()
    }

    fn reserve(&self, account_id: u64, bytes: usize) -> Result<()> {
        if bytes == 0 {
            return Ok(());
        }
        let mut state = lock_recover(&self.inner.state);
        let account = state.accounts.get(&account_id).ok_or_else(|| {
            HawDBError::Execution("query memory account is no longer registered".to_string())
        })?;
        let class = account.class;
        let owner = &account.owner;
        let account_budget = account.budget_bytes;
        let account_next = account.used_bytes.checked_add(bytes).ok_or_else(|| {
            account.account_budget_error(
                bytes,
                format!(
                    "query memory account {owner} ({}) byte accounting overflow",
                    class.as_str()
                ),
            )
        })?;
        if account_next > account_budget {
            return Err(account.account_budget_error(bytes, format!(
                "query memory account {owner} ({}) would use {account_next} bytes, exceeding its {account_budget}-byte budget",
                class.as_str()
            )));
        }
        let query_error = |message| {
            account.budget_error(
                ResourceBudgetScope::Query,
                bytes,
                state.used_bytes,
                self.inner.budget_bytes,
                self.inner.budget_option,
                message,
            )
        };
        let root_next = state.used_bytes.checked_add(bytes).ok_or_else(|| {
            query_error(format!(
                "query memory ledger byte accounting overflow while charging {owner} ({})",
                class.as_str()
            ))
        })?;
        if root_next > self.inner.budget_bytes {
            return Err(query_error(format!(
                "query memory ledger would use {root_next} bytes while charging {owner} ({}), exceeding query_memory_bytes {}",
                class.as_str(),
                self.inner.budget_bytes
            )));
        }

        // Complete allocation before changing any of the hierarchical counters.
        state.classes.entry(class).or_default();
        let account = state
            .accounts
            .get_mut(&account_id)
            .expect("validated query memory account remains registered");
        account.used_bytes = account_next;
        account.peak_bytes = account.peak_bytes.max(account_next);
        state.used_bytes = root_next;
        state.peak_bytes = state.peak_bytes.max(root_next);
        let class_state = state
            .classes
            .get_mut(&class)
            .expect("prepared query memory class remains registered");
        class_state.used_bytes = class_state.used_bytes.saturating_add(bytes);
        class_state.peak_bytes = class_state.peak_bytes.max(class_state.used_bytes);
        Ok(())
    }

    fn can_reserve(&self, account_id: u64, bytes: usize) -> bool {
        let state = lock_recover(&self.inner.state);
        let Some(account) = state.accounts.get(&account_id) else {
            return false;
        };
        let Some(account_next) = account.used_bytes.checked_add(bytes) else {
            return false;
        };
        let Some(root_next) = state.used_bytes.checked_add(bytes) else {
            return false;
        };
        account_next <= account.budget_bytes && root_next <= self.inner.budget_bytes
    }

    fn release(&self, account_id: u64, bytes: usize) {
        if bytes == 0 {
            return;
        }
        let mut state = lock_recover(&self.inner.state);
        let Some(account) = state.accounts.get_mut(&account_id) else {
            return;
        };
        let released = bytes.min(account.used_bytes);
        let class = account.class;
        account.used_bytes -= released;
        state.used_bytes = state.used_bytes.saturating_sub(released);
        if let Some(class_state) = state.classes.get_mut(&class) {
            class_state.used_bytes = class_state.used_bytes.saturating_sub(released);
        }
    }

    fn transfer(
        &self,
        source_account_id: u64,
        source_bytes: usize,
        target_account_id: u64,
        target_bytes: usize,
    ) -> Result<()> {
        if source_account_id == target_account_id {
            return Err(HawDBError::Execution(
                "query memory transfer requires distinct accounts".to_string(),
            ));
        }
        let mut state = lock_recover(&self.inner.state);
        let (source_class, source_used) = state
            .accounts
            .get(&source_account_id)
            .map(|account| (account.class, account.used_bytes))
            .ok_or_else(|| {
                HawDBError::Execution(
                    "source query memory account is no longer registered".to_string(),
                )
            })?;
        if source_bytes > source_used {
            return Err(HawDBError::Execution(format!(
                "query memory transfer tried to release {source_bytes} bytes from a source account using {source_used} bytes"
            )));
        }
        let target_account = state.accounts.get(&target_account_id).ok_or_else(|| {
            HawDBError::Execution("target query memory account is no longer registered".to_string())
        })?;
        let target_class = target_account.class;
        let target_owner = &target_account.owner;
        let target_budget = target_account.budget_bytes;
        let target_used = target_account.used_bytes;
        let target_next = target_used.checked_add(target_bytes).ok_or_else(|| {
            target_account.account_budget_error(
                target_bytes,
                format!(
                    "query memory account {target_owner} ({}) byte accounting overflow",
                    target_class.as_str()
                ),
            )
        })?;
        if target_next > target_budget {
            return Err(target_account.account_budget_error(target_bytes, format!(
                "query memory account {target_owner} ({}) would use {target_next} bytes, exceeding its {target_budget}-byte budget",
                target_class.as_str()
            )));
        }
        let root_after_release = state.used_bytes.checked_sub(source_bytes).ok_or_else(|| {
            HawDBError::Execution(
                "query memory ledger underflow during ownership transfer".to_string(),
            )
        })?;
        let query_error = |message| {
            target_account.budget_error(
                ResourceBudgetScope::Query,
                target_bytes.saturating_sub(source_bytes),
                state.used_bytes,
                self.inner.budget_bytes,
                self.inner.budget_option,
                message,
            )
        };
        let root_next = root_after_release.checked_add(target_bytes).ok_or_else(|| {
            query_error(format!(
                "query memory ledger byte accounting overflow while transferring ownership to {target_owner} ({})",
                target_class.as_str()
            ))
        })?;
        if root_next > self.inner.budget_bytes {
            return Err(query_error(format!(
                "query memory ledger would use {root_next} bytes while transferring ownership to {target_owner} ({}), exceeding query_memory_bytes {}",
                target_class.as_str(),
                self.inner.budget_bytes
            )));
        }

        state.classes.entry(target_class).or_default();
        let source = state
            .accounts
            .get_mut(&source_account_id)
            .expect("validated source query memory account remains registered");
        source.used_bytes -= source_bytes;
        let target = state
            .accounts
            .get_mut(&target_account_id)
            .expect("validated target query memory account remains registered");
        target.used_bytes = target_next;
        target.peak_bytes = target.peak_bytes.max(target_next);
        state.used_bytes = root_next;
        state.peak_bytes = state.peak_bytes.max(root_next);
        if let Some(class_state) = state.classes.get_mut(&source_class) {
            class_state.used_bytes = class_state.used_bytes.saturating_sub(source_bytes);
        }
        let class_state = state
            .classes
            .get_mut(&target_class)
            .expect("prepared target query memory class remains registered");
        class_state.used_bytes = class_state.used_bytes.saturating_add(target_bytes);
        class_state.peak_bytes = class_state.peak_bytes.max(class_state.used_bytes);
        Ok(())
    }
}

#[derive(Debug, Clone)]
pub struct QueryMemoryAccount {
    ledger: QueryMemoryLedger,
    account_id: u64,
    // A child ledger owns an up-front reservation in its parent. Every account
    // and outstanding lease keeps that reservation alive.
    backing: Option<Arc<QueryMemoryLease>>,
    retained_state: Option<Arc<QueryMemoryAccount>>,
}

impl QueryMemoryAccount {
    pub(crate) fn query_budget_bytes(&self) -> NonZeroUsize {
        self.ledger.budget_bytes()
    }

    pub(crate) fn with_retained_state(mut self, account: Self) -> Self {
        self.retained_state = Some(Arc::new(account));
        self
    }

    pub(crate) fn retained_state(&self) -> Option<&Self> {
        self.retained_state.as_deref()
    }

    /// Reserve the entire child allowance before dispatch, without charging
    /// individual child allocations to the query root a second time.
    pub(crate) fn sub_account(&self, budget: NonZeroUsize) -> Result<Self> {
        let backing = Arc::new(self.reserve(budget.get())?);
        let (class, owner) = {
            let state = lock_recover(&self.ledger.inner.state);
            let account = &state.accounts[&self.account_id];
            (account.class, Arc::clone(&account.owner))
        };
        let origin = "QueryMemoryAccount::sub_account(budget)";
        let mut child = QueryMemoryLedger::new_with_budget_option(budget, origin)
            .account_with_budget_option(class, owner, budget, origin);
        child.backing = Some(backing);
        Ok(child)
    }

    pub(crate) fn budget_bytes(&self) -> NonZeroUsize {
        NonZeroUsize::new(
            lock_recover(&self.ledger.inner.state).accounts[&self.account_id].budget_bytes,
        )
        .unwrap()
    }

    /// Remaining capacity in both this account and its shared query root.
    /// This is a sizing observation; the following reservation still arbitrates
    /// concurrent owners and can fail without advancing the producer.
    #[doc(hidden)]
    pub fn available_bytes(&self) -> usize {
        let state = lock_recover(&self.ledger.inner.state);
        let account = &state.accounts[&self.account_id];
        account.budget_bytes.saturating_sub(account.used_bytes).min(
            self.ledger
                .inner
                .budget_bytes
                .saturating_sub(state.used_bytes),
        )
    }

    pub(crate) fn peak_bytes(&self) -> usize {
        lock_recover(&self.ledger.inner.state).accounts[&self.account_id].peak_bytes
    }

    pub(crate) fn query_available_bytes(&self) -> usize {
        let state = lock_recover(&self.ledger.inner.state);
        self.ledger
            .inner
            .budget_bytes
            .saturating_sub(state.used_bytes)
    }

    pub(crate) fn sibling(
        &self,
        class: QueryMemoryClass,
        owner: impl Into<Arc<str>>,
        budget_bytes: NonZeroUsize,
    ) -> Self {
        self.sibling_with_budget_option(
            class,
            owner,
            budget_bytes,
            "QueryMemoryAccount::sibling(budget_bytes)",
        )
    }

    pub(crate) fn sibling_with_budget_option(
        &self,
        class: QueryMemoryClass,
        owner: impl Into<Arc<str>>,
        budget_bytes: NonZeroUsize,
        budget_option: &'static str,
    ) -> Self {
        let mut account =
            self.ledger
                .account_with_budget_option(class, owner, budget_bytes, budget_option);
        account.backing = self.backing.clone();
        account
    }

    pub(crate) fn can_reserve(&self, bytes: usize) -> bool {
        self.ledger.can_reserve(self.account_id, bytes)
    }

    pub fn reserve(&self, bytes: usize) -> Result<QueryMemoryLease> {
        self.ledger.reserve(self.account_id, bytes)?;
        Ok(QueryMemoryLease {
            account: self.clone(),
            bytes,
        })
    }
}

#[derive(Debug)]
pub struct QueryMemoryLease {
    account: QueryMemoryAccount,
    bytes: usize,
}

impl QueryMemoryLease {
    pub(crate) fn operator_budget_error(
        &self,
        requested: usize,
        used: usize,
        budget: usize,
        budget_option: &'static str,
        message: String,
    ) -> HawDBError {
        let state = lock_recover(&self.account.ledger.inner.state);
        state.accounts[&self.account.account_id].budget_error(
            ResourceBudgetScope::Operator,
            requested,
            used,
            budget,
            budget_option,
            message,
        )
    }

    pub(crate) fn belongs_to(&self, account: &QueryMemoryAccount) -> bool {
        self.account.account_id == account.account_id
            && Arc::ptr_eq(&self.account.ledger.inner, &account.ledger.inner)
    }

    pub fn bytes(&self) -> usize {
        self.bytes
    }

    /// Transfers part of this already admitted charge to an independent owner.
    /// The ledger charge never disappears or grows during the handoff.
    pub(crate) fn split_off(&mut self, bytes: usize) -> Result<Self> {
        if bytes > self.bytes {
            return Err(HawDBError::Execution(
                "query memory lease split exceeds its admitted charge".into(),
            ));
        }
        self.bytes -= bytes;
        Ok(Self {
            account: self.account.clone(),
            bytes,
        })
    }

    /// Takes over another lease's charge on the same account. The ledger
    /// charge is unchanged; a lease from another account is returned intact.
    pub(crate) fn absorb(&mut self, mut other: Self) -> std::result::Result<(), Self> {
        if self.account.account_id != other.account.account_id
            || !Arc::ptr_eq(&self.account.ledger.inner, &other.account.ledger.inner)
        {
            return Err(other);
        }
        self.bytes = self.bytes.saturating_add(other.bytes);
        other.bytes = 0;
        Ok(())
    }

    pub fn grow(&mut self, bytes: usize) -> Result<()> {
        self.account
            .ledger
            .reserve(self.account.account_id, bytes)?;
        self.bytes = self.bytes.saturating_add(bytes);
        Ok(())
    }

    pub fn shrink(&mut self, bytes: usize) {
        let released = bytes.min(self.bytes);
        self.account
            .ledger
            .release(self.account.account_id, released);
        self.bytes -= released;
    }

    pub fn reset(&mut self) {
        self.shrink(self.bytes);
    }

    pub(crate) fn transfer_to(
        &mut self,
        source_bytes: usize,
        target: &mut Self,
        target_bytes: usize,
    ) -> Result<()> {
        if !Arc::ptr_eq(&self.account.ledger.inner, &target.account.ledger.inner) {
            return Err(HawDBError::Execution(
                "query memory transfer requires accounts from the same ledger".to_string(),
            ));
        }
        if source_bytes > self.bytes {
            return Err(HawDBError::Execution(format!(
                "query memory transfer tried to release {source_bytes} bytes from a {}-byte lease",
                self.bytes,
            )));
        }
        self.account.ledger.transfer(
            self.account.account_id,
            source_bytes,
            target.account.account_id,
            target_bytes,
        )?;
        self.bytes -= source_bytes;
        target.bytes = target.bytes.saturating_add(target_bytes);
        Ok(())
    }
}

impl Drop for QueryMemoryLease {
    fn drop(&mut self) {
        self.account
            .ledger
            .release(self.account.account_id, self.bytes);
        self.bytes = 0;
    }
}

fn lock_recover<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
    // State never escapes the ledger. Checks and allocation precede counter updates,
    // and no caller code runs under this lock. Recovering poisoning preserves that
    // valid state; it is not a repair mechanism for arbitrary accounting corruption.
    mutex
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
}

#[cfg(test)]
#[path = "memory_ledger_hardening_tests.rs"]
mod hardening_tests;

#[cfg(test)]
#[path = "memory_ledger_budget_tests.rs"]
mod budget_tests;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn splitting_a_full_lease_preserves_charge_and_independent_drop() {
        let budget = NonZeroUsize::new(8).unwrap();
        let ledger = QueryMemoryLedger::new(budget);
        let account = ledger.account(QueryMemoryClass::ResultMaterialization, "retained", budget);
        let mut owner = account.reserve(8).unwrap();
        assert!(owner.split_off(9).is_err());
        assert_eq!(owner.bytes(), 8);
        let view = owner.split_off(3).unwrap();
        assert_eq!(owner.bytes(), 5);
        assert_eq!(view.bytes(), 3);
        assert_eq!(ledger.snapshot().used_bytes, 8);
        assert_eq!(ledger.snapshot().peak_bytes, 8);
        drop(owner);
        assert_eq!(ledger.snapshot().used_bytes, 3);
        drop(view);
        assert_eq!(ledger.snapshot().used_bytes, 0);
    }

    #[test]
    fn source_account_reuses_retained_cap_without_registering_an_unused_account() {
        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(1024).unwrap());
        let retained = ledger.account(
            QueryMemoryClass::BlockingState,
            "shared state",
            NonZeroUsize::new(512).unwrap(),
        );
        let source =
            ledger.source_account_retaining("source", ledger.budget_bytes(), retained.clone());
        assert_eq!(ledger.snapshot().account_count, 2);
        let mut grant = crate::store::admit_graph_read(&source, None, 512).unwrap();
        grant.retain_state().unwrap();
        assert!(retained.reserve(1).is_err());
        assert_eq!(ledger.snapshot().used_bytes, 512);
        drop(grant);
        let reused = retained.reserve(512).unwrap();
        assert_eq!(ledger.snapshot().account_count, 2);
        drop(reused);
        assert_eq!(ledger.snapshot().used_bytes, 0);
    }

    #[test]
    fn sub_accounts_reserve_parent_capacity_without_double_charging() {
        let budget = NonZeroUsize::new(100).unwrap();
        let ledger = QueryMemoryLedger::new(budget);
        let parent = ledger.account(QueryMemoryClass::BlockingState, "merge", budget);
        let retained = parent.reserve(10).unwrap();
        let left = parent.sub_account(NonZeroUsize::new(60).unwrap()).unwrap();
        let right = parent.sub_account(NonZeroUsize::new(30).unwrap()).unwrap();
        assert_eq!(parent.available_bytes(), 0);
        assert!(parent.sub_account(NonZeroUsize::MIN).is_err());
        let mut left_lease = left.reserve(60).unwrap();
        let right_lease = right.reserve(30).unwrap();
        assert!(left_lease.grow(1).is_err());
        assert_eq!(ledger.snapshot().used_bytes, 100);
        drop((left, right, retained));
        assert_eq!(ledger.snapshot().used_bytes, 90);
        drop(left_lease);
        assert_eq!(ledger.snapshot().used_bytes, 30);
        drop(right_lease);
        assert_eq!(ledger.snapshot().used_bytes, 0);
        assert_eq!(ledger.snapshot().peak_bytes, 100);
    }

    #[test]
    fn nested_accounts_and_siblings_keep_backing_until_last_owner_drops() {
        let budget = NonZeroUsize::new(8).unwrap();
        let ledger = QueryMemoryLedger::new(budget);
        let parent = ledger.account(QueryMemoryClass::BlockingState, "merge", budget);
        let child = parent.sub_account(budget).unwrap();
        let sibling = child.sibling(QueryMemoryClass::SpillStaging, "staging", budget);
        let nested = child.sub_account(NonZeroUsize::new(4).unwrap()).unwrap();
        assert!(sibling.reserve(5).is_err());
        let lease = nested.reserve(4).unwrap();
        drop((parent, child, nested));
        assert_eq!(ledger.snapshot().used_bytes, 8);
        drop(lease);
        assert_eq!(ledger.snapshot().used_bytes, 8);
        let sibling_lease = sibling.reserve(8).unwrap();
        drop(sibling);
        assert_eq!(ledger.snapshot().used_bytes, 8);
        drop(sibling_lease);
        assert_eq!(ledger.snapshot().used_bytes, 0);
    }

    #[test]
    fn failed_sub_account_admission_does_not_leak_root_capacity() {
        let budget = NonZeroUsize::new(100).unwrap();
        let ledger = QueryMemoryLedger::new(budget);
        let parent = ledger.account(QueryMemoryClass::BlockingState, "merge", budget);
        let other = ledger.account(QueryMemoryClass::PipelineBatch, "input", budget);
        let retained = other.reserve(70).unwrap();
        assert!(parent.sub_account(NonZeroUsize::new(31).unwrap()).is_err());
        assert_eq!(ledger.snapshot().used_bytes, 70);
        let child = parent.sub_account(NonZeroUsize::new(30).unwrap()).unwrap();
        assert!(child.reserve(31).is_err());
        drop((child, retained));
        assert_eq!(ledger.snapshot().used_bytes, 0);
    }

    #[test]
    fn sibling_accounts_share_one_root_budget() {
        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(10).unwrap());
        let left = ledger.account(
            QueryMemoryClass::BlockingState,
            "left",
            NonZeroUsize::new(10).unwrap(),
        );
        let right = ledger.account(
            QueryMemoryClass::PipelineBatch,
            "right",
            NonZeroUsize::new(10).unwrap(),
        );
        let left_lease = left.reserve(6).unwrap();
        let error = right.reserve(5).unwrap_err();

        assert!(error.to_string().contains("query_memory_bytes 10"));
        assert_eq!(ledger.snapshot().used_bytes, 6);
        drop(left_lease);
        assert_eq!(ledger.snapshot().used_bytes, 0);
    }

    #[test]
    fn capacity_preflight_checks_account_and_root_without_charging() {
        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(10).unwrap());
        let left = ledger.account(
            QueryMemoryClass::BlockingState,
            "left",
            NonZeroUsize::new(8).unwrap(),
        );
        let right = ledger.account(
            QueryMemoryClass::PipelineBatch,
            "right",
            NonZeroUsize::new(10).unwrap(),
        );

        assert!(left.can_reserve(8));
        assert_eq!(ledger.snapshot().used_bytes, 0);
        let left_lease = left.reserve(6).unwrap();
        assert!(!left.can_reserve(3));
        assert!(right.can_reserve(4));
        assert!(!right.can_reserve(5));
        assert_eq!(ledger.snapshot().used_bytes, 6);

        drop(left_lease);
        assert_eq!(ledger.snapshot().used_bytes, 0);
    }

    #[test]
    fn derived_sibling_account_keeps_the_original_query_root() {
        let budget = NonZeroUsize::new(8).unwrap();
        let ledger = QueryMemoryLedger::new(budget);
        let state = ledger.account(QueryMemoryClass::BlockingState, "merge", budget);
        let output = state.sibling(QueryMemoryClass::PipelineBatch, "output", budget);
        let state_lease = state.reserve(6).unwrap();
        assert!(output
            .reserve(3)
            .unwrap_err()
            .to_string()
            .contains("query_memory_bytes 8"));
        let output_lease = output.reserve(2).unwrap();
        assert_eq!(ledger.snapshot().used_bytes, 8);
        drop((state_lease, output_lease));
        assert_eq!(ledger.snapshot().used_bytes, 0);
    }

    #[test]
    fn ownership_transfer_does_not_require_double_root_capacity() {
        let budget = NonZeroUsize::new(8).unwrap();
        let ledger = QueryMemoryLedger::new(budget);
        let source = ledger.account(QueryMemoryClass::BlockingState, "source", budget);
        let target = ledger.account(QueryMemoryClass::PipelineBatch, "target", budget);
        let mut source_lease = source.reserve(8).unwrap();
        let mut target_lease = target.reserve(0).unwrap();

        source_lease.transfer_to(8, &mut target_lease, 8).unwrap();

        assert_eq!(source_lease.bytes(), 0);
        assert_eq!(target_lease.bytes(), 8);
        let snapshot = ledger.snapshot();
        assert_eq!(snapshot.used_bytes, 8);
        assert_eq!(snapshot.peak_bytes, 8);
        drop((source_lease, target_lease));
        assert_eq!(ledger.snapshot().used_bytes, 0);
    }

    #[test]
    fn failed_ownership_transfer_preserves_source_charge() {
        let budget = NonZeroUsize::new(8).unwrap();
        let ledger = QueryMemoryLedger::new(budget);
        let source = ledger.account(QueryMemoryClass::BlockingState, "source", budget);
        let target = ledger.account(QueryMemoryClass::PipelineBatch, "target", budget);
        let mut source_lease = source.reserve(8).unwrap();
        let mut target_lease = target.reserve(0).unwrap();

        let error = source_lease
            .transfer_to(8, &mut target_lease, 9)
            .unwrap_err();

        assert!(error.to_string().contains("9 bytes"));
        assert_eq!(source_lease.bytes(), 8);
        assert_eq!(target_lease.bytes(), 0);
        assert_eq!(ledger.snapshot().used_bytes, 8);
    }

    #[test]
    fn account_budget_is_enforced_before_root_budget() {
        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(100).unwrap());
        let account = ledger.account(
            QueryMemoryClass::BlockingState,
            "sort",
            NonZeroUsize::new(4).unwrap(),
        );

        let error = account.reserve(5).unwrap_err();

        assert!(error.to_string().contains("sort"));
        assert!(error.to_string().contains("4-byte budget"));
        assert_eq!(ledger.snapshot().used_bytes, 0);
    }

    #[test]
    fn lease_growth_shrink_and_drop_update_hierarchy() {
        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(100).unwrap());
        let account = ledger.account(
            QueryMemoryClass::MorselOutput,
            "reorder",
            NonZeroUsize::new(80).unwrap(),
        );
        let mut lease = account.reserve(20).unwrap();
        lease.grow(30).unwrap();
        lease.shrink(10);
        assert_eq!(lease.bytes(), 40);
        assert_eq!(ledger.snapshot().used_bytes, 40);

        drop(lease);
        let snapshot = ledger.snapshot();
        assert_eq!(snapshot.used_bytes, 0);
        assert_eq!(snapshot.peak_bytes, 50);
        assert_eq!(snapshot.classes[0].peak_bytes, 50);
    }

    #[test]
    fn unwind_releases_query_memory_lease() {
        let ledger = QueryMemoryLedger::new(NonZeroUsize::new(100).unwrap());
        let account = ledger.account(
            QueryMemoryClass::SpillStaging,
            "spill",
            NonZeroUsize::new(100).unwrap(),
        );
        let result = std::panic::catch_unwind(|| {
            let _lease = account.reserve(80).unwrap();
            panic!("injected panic");
        });

        assert!(result.is_err());
        assert_eq!(ledger.snapshot().used_bytes, 0);
    }
}